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

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

你可以这样做：

>>>np.diag(A[range(3),B]).reshape(B.shape)
array([[2],
       [4],
       [8]])

如果你想使用 choose，可以这样写：np.choose(B.ravel(), A.T).reshape(B.shape)。

Answer 2

你可以这样做：

>>> A[np.arange(len(A)), B.ravel()].reshape(B.shape)
array([[2],
       [4],
       [8]])

Answer 3

为了让答案更完整... 花式索引数组会被调整到一个共同的形状，所以下面的代码也能正常工作，这样你就不需要最后再调整形状了：

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