以numpy为另一个数组索引

假设我有一个具有一些任意值的matrixA ：

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

并且在A中包含元素索引的matrixB ：

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

如何从A指向的B中select值，即：

 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]])