python - index - numpy where

Numeración de fantasía indexación en múltiples dimensiones

Question

Digamos que tengo una matriz numpy A de tamaño nxmxk y otra matriz B de tamaño nxm que tiene índices de 1 a k. Quiero acceder a cada segmento nx de A usando el índice dado en este lugar en B, dándome una matriz de tamaño nx m.

Editar: ¡aparentemente eso no es lo que quiero! [[Puedo lograr esto usando la take así:

A.take(B) ]] final editar

¿Se puede lograr esto usando indexación sofisticada? Hubiera pensado que A[B] daría el mismo resultado, pero eso da como resultado una matriz de tamaño nxmxmxk (que realmente no entiendo).

La razón por la que no quiero usar take es que quiero poder asignarle algo a esta porción, como

A[B] = 1

La única solución de trabajo que tengo hasta ahora es

A.reshape(-1, k)[np.arange(n * m), B.ravel()].reshape(n, m)

pero seguramente tiene que haber una manera más fácil?

Answer 1

Suponer

import numpy as np
np.random.seed(0)
n,m,k = 2,3,5
A = np.arange(n*m*k,0,-1).reshape((n,m,k))
print(A)
# [[[30 29 28 27 26]
# [25 24 23 22 21]
# [20 19 18 17 16]]
# [[15 14 13 12 11]
# [10 9 8 7 6]
# [ 5 4 3 2 1]]]
B = np.random.randint(k, size=(n,m))
print(B)
# [[4 0 3]
# [3 3 1]]

Para crear esta matriz,

print(A.reshape(-1, k)[np.arange(n * m), B.ravel()])
# [26 25 17 12 7 4]

como una matriz nxm usando indexación elegante:

i,j = np.ogrid[0:n, 0:m]
print(A[i, j, B])
# [[26 25 17]
# [12 7 4]]