Encuentre dónde una matriz NumPy es igual a cualquier valor en una lista de valores

Tengo una matriz de enteros y quiero encontrar dónde esa matriz es igual a cualquier valor en una lista de valores múltiples.

Esto se puede hacer fácilmente tratando cada valor individualmente, o usando varias declaraciones “o” en un bucle, pero creo que debe haber una forma mejor / más rápida de hacerlo. En realidad estoy tratando con matrices de tamaño 4000 x 2000 , pero aquí hay una edición simplificada del problema:

 fake = arange(9).reshape((3,3)) array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) want = (fake==0) + (fake==2) + (fake==6) + (fake==8) print want array([[ True, False, True], [False, False, False], [ True, False, True]], dtype=bool) 

Lo que me gustaría es una forma de obtener un solo comando que incluya fake y la lista de valores [0, 2, 6, 8] .

Supongo que hay un paquete que ya tiene esto incluido que sería significativamente más rápido que si escribiera una función con un bucle en Python.

La función numpy.in1d parece hacer lo que quieres. El único problema es que solo funciona en matrices 1d, por lo que deberías usarlo así:

 In [9]: np.in1d(fake, [0,2,6,8]).reshape(fake.shape) Out[9]: array([[ True, False, True], [False, False, False], [ True, False, True]], dtype=bool) 

No tengo ni idea de por qué esto se limita a matrices 1d solamente. Mirando su código fuente , primero parece aplanar las dos matrices, después de lo cual hace algunos trucos de clasificación inteligentes. Pero nada lo detendría para que el resultado no se repita al final, como tuve que hacer a mano aquí.

La respuesta de @Bas es la que probablemente estés buscando. Pero aquí hay otra forma de hacerlo, usando el truco de vectorize de numpy:

 import numpy as np S = set([0,2,6,8]) @np.vectorize def contained(x): return x in S contained(fake) => array([[ True, False, True], [False, False, False], [ True, False, True]], dtype=bool) 

El problema de esta solución es que contained() se llama para cada elemento (es decir, en python-space), lo que hace que esto sea mucho más lento que una solución de número puro.

NumPy 0.13+

A partir de NumPy v0.13, puede usar np.isin , que funciona en matrices multidimensionales:

 >>> element = 2*np.arange(4).reshape((2, 2)) >>> element array([[0, 2], [4, 6]]) >>> test_elements = [1, 2, 4, 8] >>> mask = np.isin(element, test_elements) >>> mask array([[ False, True], [ True, False]]) 

NumPy pre-0.13

La respuesta aceptada con np.in1d solo funciona con matrices 1d y requiere una nueva np.in1d para obtener el resultado deseado. Esto es bueno para las versiones de NumPy antes v0.13.