función de estilo MATLAB find() en Python

Question

En MATLAB es fácil encontrar los índices de valores que cumplan una condición particular:

>> a = [1,2,3,1,2,3,1,2,3]; 
>> find(a > 2)  % find the indecies where this condition is true 
[3, 6, 9]   % (MATLAB uses 1-based indexing) 
>> a(find(a > 2)) % get the values at those locations 
[3, 3, 3] 

¿Cuál sería la mejor manera de hacer esto en Python?

Hasta ahora, he encontrado lo siguiente. Para obtener sólo los valores:

>>> a = [1,2,3,1,2,3,1,2,3] 
>>> [val for val in a if val > 2] 
[3, 3, 3] 

Pero si quiero que el índice de cada uno de esos valores es un poco más complicado:

>>> a = [1,2,3,1,2,3,1,2,3] 
>>> inds = [i for (i, val) in enumerate(a) if val > 2] 
>>> inds 
[2, 5, 8] 
>>> [val for (i, val) in enumerate(a) if i in inds] 
[3, 3, 3] 

¿Hay una mejor manera de hacer esto en Python, especialmente para condiciones arbitrarias (no solo 'val> 2')?

Encontré funciones equivalentes a MATLAB 'buscar' en NumPy pero actualmente no tengo acceso a esas bibliotecas.

Answer 1

Se puede hacer una función que toma un parámetro exigible que se utilizará en la parte condición de su lista por comprensión. A continuación, puede utilizar un objeto función lambda u otro para pasar su condición arbitraria:

def indices(a, func): 
    return [i for (i, val) in enumerate(a) if func(val)] 

a = [1, 2, 3, 1, 2, 3, 1, 2, 3] 

inds = indices(a, lambda x: x > 2) 

>>> inds 
[2, 5, 8] 

Es un poco más cerca de su ejemplo Matlab, sin tener que cargar todos numpy.

Answer 2

en numpy tiene where:

>> import numpy as np 
>> x = np.random.randint(0, 20, 10) 
>> x 
array([14, 13, 1, 15, 8, 0, 17, 11, 19, 13]) 
>> np.where(x > 10) 
(array([0, 1, 3, 6, 7, 8, 9], dtype=int64),) 

Answer 3

para obtener valores con condiciones arbitrarias, podría utilizar filter() con una función lambda:

>>> a = [1,2,3,1,2,3,1,2,3] 
>>> filter(lambda x: x > 2, a) 
[3, 3, 3] 

Una posible manera de obtener los índices sería utilizar enumerate() para construir una tupla con ambos índices y valores, y entonces filtro que:

>>> a = [1,2,3,1,2,3,1,2,3] 
>>> aind = tuple(enumerate(a)) 
>>> print aind 
((0, 1), (1, 2), (2, 3), (3, 1), (4, 2), (5, 3), (6, 1), (7, 2), (8, 3)) 
>>> filter(lambda x: x[1] > 2, aind) 
((2, 3), (5, 3), (8, 3)) 

Answer 4

Por qué no utilizar esto:

[i for i in range(len(a)) if a[i] > 2] 

o para condiciones arbitrarias, definir una función f para su condición y hacer:

[i for i in range(len(a)) if f(a[i])] 

Answer 5

He estado tratando de encontrar una forma rápida de hacer tal cosa exacta, y esto es lo que me topé (usos numpy para su comparación vector rápido):

a_bool = numpy.array(a) > 2 
inds = [i for (i, val) in enumerate(a_bool) if val] 

resulta que esto es mucho más rápido que:

inds = [i for (i, val) in enumerate(a) if val > 2] 

parece que Pyt hon es más rápido en comparación cuando se hace en una matriz numpy, y/o más rápido en hacer comprensiones de lista cuando solo se verifica la verdad en lugar de la comparación.

Editar:

que estaba volviendo a visitar mi código y me encontré con una posiblemente menos intensivo de memoria, un poco más rápido, y el modo super-concisa de hacer esto en una sola línea:

inds = np.arange(len(a))[ a < 2 ] 

Answer 6

O use la función no nero de numpy:

import numpy as np 
a = np.array([1,2,3,4,5]) 
inds = np.nonzero(a>2) 
a[inds] 
array([3, 4, 5]) 

Answer 7

El código de búsqueda de Matlab tiene dos argumentos. El código de John explica el primer argumento pero no el segundo. Por ejemplo, si desea saber en qué parte del índice se cumple la condición: la función de Mtlab sería:

find(x>2,1) 

Mediante el código de Juan, todo lo que tiene que hacer es añadir una [x] al final de los índices función, donde x es el número de índice que estás buscando.

def indices(a, func): 
    return [i for (i, val) in enumerate(a) if func(val)] 

a = [1, 2, 3, 1, 2, 3, 1, 2, 3] 

inds = indices(a, lambda x: x > 2)[0] #[0] being the 2nd matlab argument 

que devuelve >>> 2, el primer índice exceda 2.

Answer 8

La rutina numpy más comúnmente utilizado para esta aplicación es numpy.where(); sin embargo, creo que funciona igual que numpy.nonzero().

import numpy 
a = numpy.array([1,2,3,4,5]) 
inds = numpy.where(a>2) 

para obtener los valores, puede almacenar los índices y cortarlas en rodajas withe:

a[inds] 

o puede pasar la matriz como un parámetro opcional:

numpy.where(a>2, a) 

o múltiple matrices:

b = numpy.array([11,22,33,44,55]) 
numpy.where(a>2, a, b) 

Answer 9

Creo que pude haber encontrado un sustituto rápido y sencillo. BTW Creo que la función np.where() no es muy satisfactoria, en cierto sentido, contiene una fila molesta de elemento cero.

import matplotlib.mlab as mlab 
a = np.random.randn(1,5) 
print a 

>> [[ 1.36406736 1.45217257 -0.06896245 0.98429727 -0.59281957]] 

idx = mlab.find(a<0) 
print idx 
type(idx) 

>> [2 4] 
>> np.ndarray 

mejor, Da