Znajdowanie indeksu elementów na podstawie warunku przy użyciu funkcji rozumienia list w języku Python

119

Poniższy kod Pythona wydaje się być bardzo rozwlekły, gdy pochodzi z tła Matlab

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

Będąc w Matlabie mogę napisać:

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

Czy istnieje metoda krótkiej ręki na napisanie tego w Pythonie, czy po prostu trzymam się długiej wersji?


Dziękuję za wszystkie sugestie i wyjaśnienie uzasadnienia składni Pythona.

Po znalezieniu następujących informacji na stronie numpy myślę, że znalazłem rozwiązanie, które mi się podoba:

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

Zastosowanie informacji z tej witryny do mojego problemu powyżej dałoby:

>>> 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]

Poniższe powinno wtedy działać (ale nie mam pod ręką interpretera Pythona, aby go przetestować):

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]
Zawietrzny
źródło
6
A co powiesz [idx for idx in range(len(a)) if a[idx] > 2]? Powodem, dla którego jest to trochę niewygodne w Pythonie, jest to, że nie używa on indeksów tak często, jak inne języki.
NullUserException,

Odpowiedzi:

77
  • W Pythonie w ogóle nie używałbyś do tego indeksów, ale zająłbyś się tylko wartościami… [value for value in a if value > 2]. Zwykle obsługa indeksów oznacza, że ​​nie robisz czegoś w najlepszy sposób.

  • Jeżeli nie potrzebujemy API podobny do Matlaba, należałoby użyć numpy , pakiet dla wielowymiarowych tablic i matematyki numerycznej w Pythonie, który jest mocno inspirowany przez Matlab. Używałbyś tablicy numpy zamiast listy.

    >>> 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
źródło
2
masz listy, jedną dla zakresów i jedną dla kątów, chcesz odfiltrować wartości zakresów, które są powyżej pewnego progu. W jaki sposób można również filtrować kąty odpowiadające tym zakresom w „najlepszy sposób”?
Mehdi,
3
filtered_ranges_and_angles = [(range, angle) for range, angle in zip(ranges, angles) if should_be_kept(range)]
Mike Graham
7
„W Pythonie w ogóle byś nie używał do tego indeksów, ale po prostu zajmij się wartościami” to stwierdzenie pokazuje, że nie wykonałeś wystarczającej analizy danych i modelowania uczenia maszynowego. Indeksy jednego tensora oparte na określonym warunku są używane do filtrowania innego tensora.
horaceT,
63

Inny sposób:

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

Ogólnie rzecz biorąc, pamiętaj, że chociaż findjest to gotowa funkcja, rozumienie list składanych jest ogólnym, a zatem bardzo potężnym rozwiązaniem . Nic nie stoi na przeszkodzie, aby napisać findfunkcję w Pythonie i użyć jej później, jak chcesz. To znaczy:

>>> 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]

Zauważ, że używam tutaj list do naśladowania Matlaba. Byłoby bardziej Pythonic użycie generatorów i iteratorów.

Eli Bendersky
źródło
2
OP mógł napisać to jako [i for i,v in enumerate(a) if v > 2]zamiast tego.
NullUserException,
To nie jest krótsze, jest dłuższe. Wymień indexsię ii valueze vw oryginale i policzyć znaki.
agf
@NullUser, agf: masz rację, ale głównym punktem jest druga część :)
Eli Bendersky
1
Korzystanie z funkcji enumerateover range(len(...))jest bardziej niezawodne i wydajniejsze.
Mike Graham
1
@Mike Graham: Zgadzam się - zmienię tę find_indicesfunkcję do użyciaenumerate
Eli Bendersky
22

U mnie działa dobrze:

>>> import numpy as np
>>> a = np.array([1, 2, 3, 1, 2, 3])
>>> np.where(a > 2)[0]
[2 5]
Alexander Cyberman
źródło
6

Może inne pytanie brzmi: „co zamierzasz zrobić z tymi indeksami, kiedy już je zdobędziesz?” Jeśli masz zamiar użyć ich do utworzenia kolejnej listy, to w Pythonie są one niepotrzebnym środkowym krokiem. Jeśli chcesz, aby wszystkie wartości pasowały do ​​danego warunku, po prostu użyj wbudowanego filtra:

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

Lub napisz własną listę zawierającą:

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

Jeśli chcesz usunąć je z listy, to sposobem Pythona niekoniecznie jest usunięcie z listy, ale napisanie listy złożonej tak, jakbyś tworzył nową listę i przypisując ją z powrotem w miejscu za pomocą listvar[:]po lewej stronie -bok:

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

Matlab dostarcza, findponieważ jego model zorientowany na tablice działa poprzez wybieranie elementów za pomocą ich indeksów tablicowych. Oczywiście możesz to zrobić w Pythonie, ale bardziej Pythonowym sposobem jest używanie iteratorów i generatorów, jak już wspomniał @EliBendersky.

PaulMcG
źródło
Paul, jeszcze nie znalazłem takiej potrzeby w skrypcie / funkcji / klasie. Bardziej chodzi o interaktywne testowanie klasy, którą piszę.
Lee
@Mike - dzięki za edycję, ale naprawdę miałem na myśli a[:] = ...- zobacz odpowiedź Alexa Martelliego na to pytanie stackoverflow.com/questions/1352885/… .
PaulMcG
@Paul, założyłem (i miałem nadzieję!), Że tak naprawdę nie miałeś na myśli tego z opisu, że zamierzasz „utworzyć nową listę”; Uważam, że programy są łatwiejsze do zrozumienia i utrzymania, gdy bardzo oszczędnie mutują istniejące dane. W każdym razie przepraszam, że przekroczyłem granicę - z pewnością powinieneś móc edytować swój post z powrotem do tego, co chcesz.
Mike Graham
6

Nawet jeśli jest to późna odpowiedź: myślę, że to nadal bardzo dobre pytanie, a IMHO Python (bez dodatkowych bibliotek lub zestawów narzędzi, takich jak numpy) wciąż brakuje wygodnej metody dostępu do indeksów elementów listy według ręcznie zdefiniowanego filtru.

Możesz ręcznie zdefiniować funkcję, która zapewnia taką funkcjonalność:

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))

Wydajność: [0, 4]

W mojej wyobraźni idealnym sposobem byłoby stworzenie klasy potomnej listy i dodanie indeksów jako metody klasowej. W ten sposób potrzebna byłaby tylko metoda filtrowania:

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)

Omówiłem nieco więcej na ten temat tutaj: http://tinyurl.com/jajrr87

Gerhard Hagerer
źródło