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

Question

说我有一个数组

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 ，但我没有做到这一点。

Answer 1

为了答案的完整性......花式索引数组被广播到一个共同的形状，所以以下也有效，并且让你免于最后的 reshape :

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