Индексирование одного массива другим в 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]]

Ответ 1

Вы можете использовать NumPy 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]])

Ответ 2

Ниже приведено решение с использованием цикла 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]]