promedio - que es un array en python

Seleccionando filas de un NumPy ndarray

Question

Quiero seleccionar solo ciertas filas de una matriz NumPy función del valor en la segunda columna. Por ejemplo, esta matriz de prueba tiene enteros de 1 a 10 en la segunda columna.

>>> test = numpy.array([numpy.arange(100), numpy.random.randint(1, 11, 100)]).transpose()
>>> test[:10, :]
array([[ 0,  6],
       [ 1,  7],
       [ 2, 10],
       [ 3,  4],
       [ 4,  1],
       [ 5, 10],
       [ 6,  6],
       [ 7,  4],
       [ 8,  6],
       [ 9,  7]])

Si quisiera solo filas donde el segundo valor es 4, es fácil:

>>> test[test[:, 1] == 4]
array([[ 3,  4],
       [ 7,  4],
       [16,  4],
       ...
       [81,  4],
       [83,  4],
       [88,  4]])

Pero, ¿cómo logro el mismo resultado cuando hay más de un valor deseado?

La lista deseada puede ser de longitud arbitraria. Por ejemplo, es posible que desee que todas las filas donde la segunda columna sea 2, 4 o 6:

>>> wanted = [2, 4, 6]

La única forma en que se me ha ocurrido es usar la comprensión de lista y luego convertirla de nuevo en una matriz y parece demasiado complicada, aunque funciona:

>>> test[numpy.array([test[x, 1] in wanted for x in range(len(test))])]
array([[ 0,  6],
       [ 3,  4],
       [ 6,  6],
       ...
       [90,  2],
       [91,  6],
       [92,  2]])

¿Hay una mejor manera de hacer esto en NumPy que me estoy perdiendo?

Answer 1

Esto es dos veces más rápido que la variante de Amnon para len (prueba) = 1000:

wanted = (2,4,6)
wanted2 = numpy.expand_dims(wanted, 1)
print test[numpy.any(test[:, 1] == wanted2, 0), :]

Answer 2

La siguiente solución debería ser más rápida que la solución de Amnon, ya que la wanted aumenta:

# Much faster look up than with lists, for larger lists:
wanted_set = set(wanted)

@numpy.vectorize
def selected(elmt): return elmt in wanted_set
# Or: selected = numpy.vectorize(wanted_set.__contains__)

print test[selected(test[:, 1])]

De hecho, tiene la ventaja de buscar a través de la matriz de test solo una vez (en lugar de tantas como len(wanted) veces como en la respuesta de Amnon). También utiliza la búsqueda rápida de elementos incorporada de Python en conjuntos , que son mucho más rápidos para esto que las listas. También es rápido porque utiliza los bucles rápidos de Numpy. También obtiene la optimización del operador in : una vez que un elemento wanted coincide, los elementos restantes no tienen que ser probados (a diferencia del enfoque "lógico o" de Amnon, donde todos los elementos wanted se prueban sin importar qué) .

Alternativamente, puede usar la siguiente frase de una línea, que también pasa por su matriz solo una vez:

test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)]

Sin embargo, esto es mucho más lento, ya que extrae el elemento en la segunda columna en cada iteración (en lugar de hacerlo en una pasada, como en la primera solución de esta respuesta).

Answer 3

numpy.in1d es lo que estás buscando:

print test[numpy.in1d(test[:,1], wanted)]

Fácilmente debería ser la solución más rápida si se desea que sea grande; Además, es el más legible, digo id.

Answer 4

test[numpy.logical_or.reduce([test[:,1] == x for x in wanted])]

El resultado debería ser más rápido que la versión original, ya que NumPy hace los bucles internos en lugar de Python.