Encontrar el índice de elementos basado en una condición usando la comprensión de la lista de Python

119

El siguiente código de Python parece ser muy largo cuando proviene de un fondo de Matlab

>>> a = [1, 2, 3, 1, 2, 3]
>>> [index for index,value in enumerate(a) if value > 2]
[2, 5]

Cuando en Matlab puedo escribir:

>> a = [1, 2, 3, 1, 2, 3];
>> find(a>2)
ans =
     3     6

¿Existe un método de mano corta para escribir esto en Python, o simplemente me quedo con la versión larga?


Gracias por todas las sugerencias y explicaciones del fundamento de la sintaxis de Python.

Después de encontrar lo siguiente en el sitio web de numpy, creo que encontré una solución que me gusta:

http://docs.scipy.org/doc/numpy/user/basics.indexing.html#boolean-or-mask-index-arrays

Aplicar la información de ese sitio web a mi problema anterior, daría lo siguiente:

>>> from numpy import array
>>> a = array([1, 2, 3, 1, 2, 3])
>>> b = a>2 
array([False, False, True, False, False, True], dtype=bool)
>>> r = array(range(len(b)))
>>> r(b)
[2, 5]

Lo siguiente debería funcionar (pero no tengo un intérprete de Python a mano para probarlo):

class my_array(numpy.array):
    def find(self, b):
        r = array(range(len(b)))
        return r(b)


>>> a = my_array([1, 2, 3, 1, 2, 3])
>>> a.find(a>2)
[2, 5]
Sotavento
fuente
6
¿Qué tal [idx for idx in range(len(a)) if a[idx] > 2]? La razón por la que esto es un poco incómodo de hacer en Python es porque no usa índices tanto como otros lenguajes.
NullUserException

Respuestas:

77
  • En Python, no usarías índices para esto en absoluto, solo lidiarías con los valores— [value for value in a if value > 2]. Por lo general, tratar con índices significa que no está haciendo algo de la mejor manera.

  • Si usted no necesita una API similar a Matlab, se usaría numpy , un paquete de matrices multidimensionales y matemáticas numérica en Python que está fuertemente inspirado en Matlab. Estaría usando una matriz numpy en lugar de una lista.

    >>> import numpy
    >>> a = numpy.array([1, 2, 3, 1, 2, 3])
    >>> a
    array([1, 2, 3, 1, 2, 3])
    >>> numpy.where(a > 2)
    (array([2, 5]),)
    >>> a > 2
    array([False, False,  True, False, False,  True], dtype=bool)
    >>> a[numpy.where(a > 2)]
    array([3, 3])
    >>> a[a > 2]
    array([3, 3])
Mike Graham
fuente
2
tiene listas, una para rangos y otra para ángulos, desea filtrar los valores de rango que están por encima de algún umbral. ¿Cómo se filtran también los ángulos correspondientes a esos rangos de la "mejor manera"?
Mehdi
3
filtered_ranges_and_angles = [(range, angle) for range, angle in zip(ranges, angles) if should_be_kept(range)]
Mike Graham
7
"En Python, no usarías índices para esto en absoluto, sino que simplemente lidias con los valores", esta declaración muestra que no has realizado suficientes análisis de datos y modelado de aprendizaje automático. Los índices de un tensor basados ​​en cierta condición se utilizan para filtrar otro tensor.
horaceT
63

De otra manera:

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

En general, recuerde que si bien findes una función lista para usar, las listas por comprensión son una solución general y, por lo tanto, muy poderosa . Nada le impide escribir una findfunción en Python y usarla más tarde como desee. Es decir:

>>> def find_indices(lst, condition):
...   return [i for i, elem in enumerate(lst) if condition(elem)]
... 
>>> find_indices(a, lambda e: e > 2)
[2, 5]

Tenga en cuenta que estoy usando listas aquí para imitar Matlab. Sería más Pythonic usar generadores e iteradores.

Eli Bendersky
fuente
2
El OP podría haberlo escrito como en su [i for i,v in enumerate(a) if v > 2]lugar.
NullUserException
Eso no es más corto, es más largo. Reemplace indexcon iy valuecon ven el original y cuente los caracteres.
agf
@NullUser, agf: tienes razón, pero el punto principal es la segunda parte :)
Eli Bendersky
1
Usar enumerateover range(len(...))es más robusto y más eficiente.
Mike Graham
1
@Mike Graham: Estoy de acuerdo, cambiaré la find_indicesfunción para usarenumerate
Eli Bendersky
22

Para mi funciona bien:

>>> import numpy as np
>>> a = np.array([1, 2, 3, 1, 2, 3])
>>> np.where(a > 2)[0]
[2 5]
Alexander Cyberman
fuente
6

Quizás otra pregunta sea, "¿qué vas a hacer con esos índices una vez que los obtengas?" Si los va a usar para crear otra lista, en Python, son un paso intermedio innecesario. Si desea que todos los valores coincidan con una condición determinada, simplemente use el filtro incorporado:

matchingVals = filter(lambda x : x>2, a)

O escriba su propia lista de comprensión:

matchingVals = [x for x in a if x > 2]

Si desea eliminarlos de la lista, entonces la forma Pythonic no es necesariamente eliminarlos de la lista, sino escribir una lista de comprensión como si estuviera creando una nueva lista y volviendo a asignar en el lugar usando el listvar[:]de la izquierda. -lado:

a[:] = [x for x in a if x <= 2]

Matlab suministra findporque su modelo centrado en matrices funciona seleccionando elementos utilizando sus índices de matriz. Usted puede hacer esto en Python, sin duda, pero la forma más Pythonic está utilizando iteradores y generadores, como ya se ha mencionado por @EliBendersky.

PaulMcG
fuente
Paul, todavía no he encontrado la necesidad de esto en un script / función / clase. Es más para pruebas interactivas de una clase que estoy escribiendo.
Lee
@Mike, gracias por la edición, pero realmente quise decir a[:] = ..., mira la respuesta de Alex Martelli a esta pregunta stackoverflow.com/questions/1352885/… .
PaulMcG
@Paul, asumí (¡y esperaba!) Que realmente no lo decía en serio a partir de su descripción que iba a "crear una nueva lista"; Encuentro que los programas tienden a ser más fáciles de entender y mantener cuando mutan datos existentes con mucha moderación. En cualquier caso, lamento excederme, sin duda debería poder editar su publicación de nuevo a lo que desee.
Mike Graham
6

Incluso si es una respuesta tardía: creo que esta sigue siendo una muy buena pregunta y, en mi humilde opinión, Python (sin bibliotecas o kits de herramientas adicionales como numpy) todavía carece de un método conveniente para acceder a los índices de los elementos de la lista de acuerdo con un filtro definido manualmente.

Puede definir manualmente una función, que proporciona esa funcionalidad:

def indices(list, filtr=lambda x: bool(x)):
    return [i for i,x in enumerate(list) if filtr(x)]

print(indices([1,0,3,5,1], lambda x: x==1))

Rendimientos: [0, 4]

En mi imaginación, la manera perfecta sería hacer una lista de clase secundaria y agregar la función de índices como método de clase. De esta manera solo se necesitaría el método de filtro:

class MyList(list):
    def __init__(self, *args):
        list.__init__(self, *args)
    def indices(self, filtr=lambda x: bool(x)):
        return [i for i,x in enumerate(self) if filtr(x)]

my_list = MyList([1,0,3,5,1])
my_list.indices(lambda x: x==1)

Elaboré un poco más sobre ese tema aquí: http://tinyurl.com/jajrr87

Gerhard Hagerer
fuente