Numpy数组行通过索引数组索引

Question

想象一下，我有以下（3,2）Numpy数组A

A = np.array([[1, 2],
       [3, 4],
       [5, 6]])

并且我希望通过numpy Array B中可用的索引来逐列索引此数组：

B = np.array([[1,0],[2,0]])

所以我想在第0列中获取第1行和第0行，在第1列中获取第2行和第0行以获取：

C = np.array([[3,1],[6,2]])

这样做的有效方法是什么？

Answer 1

您可以使用np.arange(A.shape[1])构建列索引，转置B，以便正确地使用列索引进行广播，然后使用advanced indexing提取元素：

A = np.array([[1, 2],
       [3, 4],
       [5, 6]])

B = np.array([[1,0],[2,0]])

A[B.T, np.arange(A.shape[1])].T
#array([[3, 1],
#       [6, 2]])

行索引：

B.T
#           V second column row index
#array([[1, 2],
#       [0, 0]])
#        ^ first column row index

列索引：

np.arange(A.shape[1])
# array([0, 1])