NumPy seleccionando un índice de columna específico por fila usando una lista de índices

90

Estoy luchando por seleccionar las columnas específicas por fila de una NumPymatriz.

Supongamos que tengo la siguiente matriz a la que llamaría X:

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]

También tengo un listíndice de columnas por cada fila que llamaría Y:

[1, 0, 2]

Necesito obtener los valores:

[2]
[4]
[9]

En lugar de listcon índices Y, también puedo producir una matriz con la misma forma Xen la que cada columna tiene un valor bool/ inten el rango 0-1, lo que indica si esta es la columna requerida.

[0, 1, 0]
[1, 0, 0]
[0, 0, 1]

Sé que esto se puede hacer iterando sobre la matriz y seleccionando los valores de columna que necesito. Sin embargo, esto se ejecutará con frecuencia en grandes conjuntos de datos y es por eso que tiene que ejecutarse tan rápido como sea posible.

Por tanto, me preguntaba si hay una solución mejor.

Gracias.

Zee
fuente
¿Es la respuesta mejor para ti? stackoverflow.com/a/17081678/5046896
GoingMyWay

Respuestas:

102

Si tiene una matriz booleana, puede hacer una selección directa basada en eso así:

>>> a = np.array([True, True, True, False, False])
>>> b = np.array([1,2,3,4,5])
>>> b[a]
array([1, 2, 3])

Para seguir su ejemplo inicial, puede hacer lo siguiente:

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> b = np.array([[False,True,False],[True,False,False],[False,False,True]])
>>> a[b]
array([2, 4, 9])

También puede agregar arangey hacer una selección directa sobre eso, aunque dependiendo de cómo esté generando su matriz booleana y cómo se ve su código YMMV.

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> a[np.arange(len(a)), [1,0,2]]
array([2, 4, 9])

Espero que te ayude, avísame si tienes más preguntas.

Slater Victoroff
fuente
11
+1 para el ejemplo usando arange. Esto fue particularmente útil para mí para recuperar diferentes bloques de múltiples matrices (básicamente el caso 3D de este ejemplo)
Griddo
1
Hola, ¿podrías explicar por qué tenemos que usar en arangelugar de :? Sé que tu forma funciona y la mía no, pero me gustaría entender por qué.
marcotama
@tamzord porque es una matriz numerosa y no una lista de Python de vainilla, por lo que la :sintaxis no funciona de la misma manera.
Slater Victoroff
1
@SlaterTyranus, gracias por responder. Tengo entendido, después de leer un poco, que mezclar :con la indexación avanzada significa: "para cada subespacio :, aplique la indexación avanzada dada". ¿Es correcto mi entendimiento?
marcotama
@tamzord explica lo que quieres decir con "
subespacio
35

Puedes hacer algo como esto:

In [7]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [8]: lst = [1, 0, 2]

In [9]: a[np.arange(len(a)), lst]
Out[9]: array([2, 4, 9])

Más información sobre la indexación de matrices multidimensionales: http://docs.scipy.org/doc/numpy/user/basics.indexing.html#indexing-multi-dimensional-arrays

Ashwini Chaudhary
fuente
1
luchando por entender por qué se necesita el rango en lugar de simplemente ':' o rango.
MadmanLee
@MadmanLee Hola, el uso :generará múltiples len(a)veces los resultados, en cambio, indicando que el índice de cada fila imprimirá los resultados anticipados.
GoingMyWay
1
Creo que esta es exactamente la forma correcta y elegante de resolver este problema.
GoingMyWay
6

Una forma sencilla podría verse así:

In [1]: a = np.array([[1, 2, 3],
   ...: [4, 5, 6],
   ...: [7, 8, 9]])

In [2]: y = [1, 0, 2]  #list of indices we want to select from matrix 'a'

range(a.shape[0]) volverá array([0, 1, 2])

In [3]: a[range(a.shape[0]), y] #we're selecting y indices from every row
Out[3]: array([2, 4, 9])
Dhaval Mayatra
fuente
1
Por favor, considere agregar explicaciones.
souki
@souki He añadido una explicación ahora. Gracias
Dhaval Mayatra
6

Las numpyversiones recientes han agregado un take_along_axis(y put_along_axis) que hace esta indexación limpiamente.

In [101]: a = np.arange(1,10).reshape(3,3)                                                             
In [102]: b = np.array([1,0,2])                                                                        
In [103]: np.take_along_axis(a, b[:,None], axis=1)                                                     
Out[103]: 
array([[2],
       [4],
       [9]])

Funciona de la misma forma que:

In [104]: a[np.arange(3), b]                                                                           
Out[104]: array([2, 4, 9])

pero con manejo de eje diferente. Está especialmente dirigido a aplicar los resultados de argsorty argmax.

hpaulj
fuente
3

Puede hacerlo usando iterador. Me gusta esto:

np.fromiter((row[index] for row, index in zip(X, Y)), dtype=int)

Hora:

N = 1000
X = np.zeros(shape=(N, N))
Y = np.arange(N)

#@Aशwini चhaudhary
%timeit X[np.arange(len(X)), Y]
10000 loops, best of 3: 30.7 us per loop

#mine
%timeit np.fromiter((row[index] for row, index in zip(X, Y)), dtype=int)
1000 loops, best of 3: 1.15 ms per loop

#mine
%timeit np.diag(X.T[Y])
10 loops, best of 3: 20.8 ms per loop
Kei Minagawa
fuente
1
OP mencionó que debería ejecutarse rápido en matrices grandes , por lo que sus puntos de referencia no son muy representativos. ¡Tengo curiosidad por saber cómo funciona su último método para matrices (mucho) más grandes!
@moarningsun: actualizado. np.diag(X.T[Y])es tan lento ... Pero np.diag(X.T)es tan rápido (10us). No sé por qué.
Kei Minagawa
0

Otra forma inteligente es transponer primero la matriz y luego indexarla. Finalmente, tome la diagonal, siempre es la respuesta correcta.

X = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
Y = np.array([1, 0, 2, 2])

np.diag(X.T[Y])

Paso a paso:

Matrices originales:

>>> X
array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

>>> Y
array([1, 0, 2, 2])

Transponer para que sea posible indexarlo correctamente.

>>> X.T
array([[ 1,  4,  7, 10],
       [ 2,  5,  8, 11],
       [ 3,  6,  9, 12]])

Obtenga filas en el orden Y.

>>> X.T[Y]
array([[ 2,  5,  8, 11],
       [ 1,  4,  7, 10],
       [ 3,  6,  9, 12],
       [ 3,  6,  9, 12]])

La diagonal ahora debería quedar clara.

>>> np.diag(X.T[Y])
array([ 2,  4,  9, 12]
Thomas Devoogdt
fuente
1
Esto técnicamente funciona y se ve muy elegante. Sin embargo, encuentro que este enfoque explota por completo cuando se trata de matrices grandes. En mi caso, NumPy se tragó 30 GB de intercambio y llenó mi SSD. En su lugar, recomiendo utilizar el enfoque de indexación avanzada.
5nefarious