在numpy中将一个数组索引为另一个数组

Question

假设我有一个带有一些任意值的矩阵A:

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

并且矩阵B包含A中元素的索引:

array([[0, 0, 1, 2],
       [0, 3, 2, 1],
       [3, 2, 1, 0]])

我该如何选择值一个被指出乙,即:

A[B] = [[2, 2, 4, 5],
        [1, 9, 8, 6],
        [2, 0, 7, 8]]

Answer 1

你可以用NumPy's advanced indexing-

A[np.arange(A.shape[0])[:,None],B]

人们也可以使用linear indexing-

m,n = A.shape
out = np.take(A,B + n*np.arange(m)[:,None])

样品运行 -

In [40]: A
Out[40]: 
array([[2, 4, 5, 3],
       [1, 6, 8, 9],
       [8, 7, 0, 2]])

In [41]: B
Out[41]: 
array([[0, 0, 1, 2],
       [0, 3, 2, 1],
       [3, 2, 1, 0]])

In [42]: A[np.arange(A.shape[0])[:,None],B]
Out[42]: 
array([[2, 2, 4, 5],
       [1, 9, 8, 6],
       [2, 0, 7, 8]])

In [43]: m,n = A.shape

In [44]: np.take(A,B + n*np.arange(m)[:,None])
Out[44]: 
array([[2, 2, 4, 5],
       [1, 9, 8, 6],
       [2, 0, 7, 8]])

Answer 2

更新的版本添加了一个take_along_axis可以完成这项工作的函数：

A = np.array([[ 2, 4, 5, 3], 
              [ 1, 6, 8, 9], 
              [ 8, 7, 0, 2]])  
                                          
B = np.array([[0, 0, 1, 2], 
              [0, 3, 2, 1], 
              [3, 2, 1, 0]])     
                                        
np.take_along_axis(A, B, 1)  
                                       
Out[]: 
array([[2, 2, 4, 5],
       [1, 9, 8, 6],
       [2, 0, 7, 8]])

还有一个put_along_axis.

Answer 3

我知道这是一个老问题，但使用索引的另一种方法是：

A[np.indices(B.shape)[0], B]

输出：

[[2 2 4 5]
 [1 9 8 6]
 [2 0 7 8]]