python - Indexación de una matriz por otra en numpy

Question

Supongamos que tengo una matriz A con algunos valores arbitrarios:

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

Y una matriz B que contiene índices de elementos en A:

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

¿Cómo selecciono valores de A señalados por B , es decir:

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

Answer 1

La siguiente es la solución usando for loop:

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)

Lo anterior también se puede escribir en una forma más breve de comprensión de la lista:

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)

Salida:

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

Answer 2

Las versiones más recientes han agregado una función take_along_axis que hace el trabajo:

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

También hay un put_along_axis .

Answer 3

Puede usar NumPy''s advanced indexing :

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

También se puede usar linear indexing :

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

Ejecución de muestra:

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