python - 使用另一个数组作为索引从数组中选择元素

说我有一个数组

A = array([[1,2,3],
           [4,5,6],
           [7,8,9]])

索引数组为

B = array([[1], # want [0, 1] element of A
           [0],  # want [1, 0], element of A
           [1]])  # want [2, 1] elemtn of A

通过此索引数组B，我想要一个3-by-1数组，其元素从数组A中获取，即

C = array([[2],
           [4],
           [8]])

我尝试了numpy.choose，但失败了。

最佳答案

为了使答案更完整...花式索引数组将广播为通用形状，因此以下内容也可以使用，并为您节省了最后的重塑:

>>> A[np.arange(3)[:, None], B]
array([[2],
       [4],
       [8]])