Seleccionar filas y columnas específicas de la matriz NumPy

97

Me he estado volviendo loco tratando de averiguar qué estupidez estoy haciendo mal aquí.

Estoy usando NumPy y tengo índices de fila específicos e índices de columna específicos entre los que quiero seleccionar. Aquí está la esencia de mi problema:

import numpy as np

a = np.arange(20).reshape((5,4))
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11],
#        [12, 13, 14, 15],
#        [16, 17, 18, 19]])

# If I select certain rows, it works
print a[[0, 1, 3], :]
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [12, 13, 14, 15]])

# If I select certain rows and a single column, it works
print a[[0, 1, 3], 2]
# array([ 2,  6, 14])

# But if I select certain rows AND certain columns, it fails
print a[[0,1,3], [0,2]]
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
# ValueError: shape mismatch: objects cannot be broadcast to a single shape

¿Por qué está pasando esto? Seguramente debería poder seleccionar las filas 1, 2 y 4, y las columnas 1 y 3. El resultado que espero es:

a[[0,1,3], [0,2]] => [[0,  2],
                      [4,  6],
                      [12, 14]]
python arrays numpy multidimensional-array numpy-slicing Mike C
fuente
Etiquetado numpy-slicing para mejorar la búsqueda. (Además, los términos 'corte' y 'corte' no aparecen en el texto sin formato, podríamos usar algunos duplicados con esos términos cerrados en esto)
smci
Este es un duplicado de stackoverflow.com/questions/19161512/numpy-extract-submatrix
David John Coleman II

Respuestas:

86

La indexación elegante requiere que proporcione todos los índices para cada dimensión. Está proporcionando 3 índices para el primero y solo 2 para el segundo, de ahí el error. Quieres hacer algo como esto:

>>> a[[[0, 0], [1, 1], [3, 3]], [[0,2], [0,2], [0, 2]]]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Por supuesto, es un dolor de cabeza escribir, por lo que puede dejar que la transmisión lo ayude:

>>> a[[[0], [1], [3]], [0, 2]]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Esto es mucho más sencillo de hacer si indexa con matrices, no con listas:

>>> row_idx = np.array([0, 1, 3])
>>> col_idx = np.array([0, 2])
>>> a[row_idx[:, None], col_idx]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])
Jaime
fuente
4
Gracias, ¡no sabía que podías hacer esto! La transmisión es extraña y maravillosa ... Después de dos años de numpy, todavía me estoy acostumbrando.
Praveen
2
¡Gracias! Si bien las otras respuestas respondieron mi pregunta correctamente en términos de devolver la matriz seleccionada, esta respuesta abordó eso al mismo tiempo que abordaba el problema de la asignación (cómo establecer un [[0,1,3], [0,2]] = 0 , por ejemplo).
Mike C
1
@Jaime - Ayer mismo descubrí un one-liner incorporado para hacer exactamente el truco de transmisión que sugieres: np.ix_
Praveen
1
¿Alguien podría proporcionar una explicación de por qué la sintaxis funciona así? ¿Cuál es la razón por la que funciona para los dos primeros ejemplos pero no para el tercero? Y también, ¿cómo resuelve esto encapsular los índices buscados en sus propias listas? Gracias
Aetos
2
¿Por qué las filas necesitan estar anidadas y las columnas no?
AturSams
86

Como sugiere Toan, un truco simple sería simplemente seleccionar las filas primero y luego seleccionar las columnas sobre eso .

>>> a[[0,1,3], :]            # Returns the rows you want
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [12, 13, 14, 15]])
>>> a[[0,1,3], :][:, [0,2]]  # Selects the columns you want as well
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

[Editar] El método integrado: np.ix_

Recientemente descubrí que numpy le brinda una fórmula incorporada para hacer exactamente lo que sugirió @Jaime, pero sin tener que usar la sintaxis de transmisión (que adolece de falta de legibilidad). De los documentos:

Usando ix_, uno puede construir rápidamente matrices de índices que indexarán el producto cruzado. a[np.ix_([1,3],[2,5])]devuelve la matriz [[a[1,2] a[1,5]], [a[3,2] a[3,5]]].

Entonces lo usas así:

>>> a = np.arange(20).reshape((5,4))
>>> a[np.ix_([0,1,3], [0,2])]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Y la forma en que funciona es que se encarga de alinear las matrices como sugirió Jaime, para que la transmisión se realice correctamente:

>>> np.ix_([0,1,3], [0,2])
(array([[0],
        [1],
        [3]]), array([[0, 2]]))

Además, como dice MikeC en un comentario, np.ix_tiene la ventaja de devolver una vista, que mi primera respuesta (previa a la edición) no lo hizo. Esto significa que ahora puede asignar a la matriz indexada:

>>> a[np.ix_([0,1,3], [0,2])] = -1
>>> a    
array([[-1,  1, -1,  3],
       [-1,  5, -1,  7],
       [ 8,  9, 10, 11],
       [-1, 13, -1, 15],
       [16, 17, 18, 19]])
Praveen
fuente
4
En algunas pruebas, también descubrí np.ix_que era más rápido que el método de seleccionar las primeras columnas y luego las filas (generalmente aproximadamente 2 veces más rápido en mis pruebas de matrices cuadradas de tamaños 1K-10K donde reindexa todas las filas y columnas).
Nathan
7

UTILIZAR:

 >>> a[[0,1,3]][:,[0,2]]
array([[ 0,  2],
   [ 4,  6],
   [12, 14]])

O:

>>> a[[0,1,3],::2]
array([[ 0,  2],
   [ 4,  6],
   [12, 14]])
Toan Nguyen
fuente
10
Si bien esto es correcto, debería considerar publicar un poco más de información que explique por qué es correcto.
ebarr
2

Usar np.ix_es la forma más conveniente de hacerlo (como respondieron otros), pero aquí hay otra forma interesante de hacerlo:

>>> rows = [0, 1, 3]
>>> cols = [0, 2]

>>> a[rows].T[cols].T

array([[ 0,  2],
       [ 4,  6],
       [12, 14]])
Andreas K.
fuente