在numpy中索引一个数组又一个数组

Question

在numpy中索引一个数组又一个数组

假设我有一个矩阵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，即：

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]: AOut[40]: array([[2, 4, 5, 3],        [1, 6, 8, 9],        [8, 7, 0, 2]])In [41]: BOut[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.shapeIn [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做此工作的函数：In [203]: A = np.array([[ 2, 4, 5, 3],       ...:        [ 1, 6, 8, 9],       ...:        [ 8, 7, 0, 2]])                                                      In [204]: B = np.array([[0, 0, 1, 2],       ...:        [0, 3, 2, 1],       ...:        [3, 2, 1, 0]])                                                       In [205]: np.take_along_axis(A,B,1)                                                   Out[205]: array([[2, 2, 4, 5],        [1, 9, 8, 6],        [2, 0, 7, 8]])还有一个put_along_axis.

Answer 3

以下是使用for循环：outlist = []for i in range(len(B)):     lst = []         for j in range(len(B[i])):         lst.append(A[i][B[i][j]])     outlist.append(lst)outarray = np.asarray(outlist)print(outarray)以上也可以用更简洁的列表理解形式写成：outlist = [ [A[i][B[i][j]] for j in range(len(B[i]))]                 for i in range(len(B))  ]outarray = np.asarray(outlist)print(outarray)产出：[[2 2 4 5]  [1 9 8 6]  [2 0 7 8]]