Filtrando uma lista com base em uma lista de booleanos


127

Eu tenho uma lista de valores que eu preciso filtrar dados os valores em uma lista de booleanos:

list_a = [1, 2, 4, 6]
filter = [True, False, True, False]

Eu gero uma nova lista filtrada com a seguinte linha:

filtered_list = [i for indx,i in enumerate(list_a) if filter[indx] == True]

o que resulta em:

print filtered_list
[1,4]

A linha funciona, mas parece (para mim) um pouco exagerada e eu queria saber se havia uma maneira mais simples de conseguir o mesmo.


Conselhos

Resumo de dois bons conselhos dados nas respostas abaixo:

1- Não nomeie uma lista filtercomo eu fiz, porque é uma função interna.

2- Não compare as coisas Truecomo eu fiz, if filter[idx]==True..pois é desnecessário. Basta usar if filter[idx]é suficiente.


3
Apenas para sua informação, esta é uma primitiva de computação paralela comum chamada compactação de fluxo . (É chamado de 'primitivo' não porque é simples, mas porque é usado como um bloco de construção para muitos outros algoritmos paralelos)
BlueRaja - Danny Pflughoeft 6/13

2
Algumas notas de estilo: if filter[indx] == TrueVocê não usar ==se você quiser verificar se há identidade com True, uso is. De qualquer forma, neste caso, toda a comparação é inútil, você pode simplesmente usar if filter[indx]. Por fim: nunca use o nome de um built-in como um nome de variável / módulo (estou me referindo ao nome filter). Usando algo como included, para que ifleia bem ( if included[indx]).
Bakuriu

Respostas:


184

Você está procurando itertools.compress:

>>> from itertools import compress
>>> list_a = [1, 2, 4, 6]
>>> fil = [True, False, True, False]
>>> list(compress(list_a, fil))
[1, 4]

Comparações de tempo (py3.x):

>>> list_a = [1, 2, 4, 6]
>>> fil = [True, False, True, False]
>>> %timeit list(compress(list_a, fil))
100000 loops, best of 3: 2.58 us per loop
>>> %timeit [i for (i, v) in zip(list_a, fil) if v]  #winner
100000 loops, best of 3: 1.98 us per loop

>>> list_a = [1, 2, 4, 6]*100
>>> fil = [True, False, True, False]*100
>>> %timeit list(compress(list_a, fil))              #winner
10000 loops, best of 3: 24.3 us per loop
>>> %timeit [i for (i, v) in zip(list_a, fil) if v]
10000 loops, best of 3: 82 us per loop

>>> list_a = [1, 2, 4, 6]*10000
>>> fil = [True, False, True, False]*10000
>>> %timeit list(compress(list_a, fil))              #winner
1000 loops, best of 3: 1.66 ms per loop
>>> %timeit [i for (i, v) in zip(list_a, fil) if v] 
100 loops, best of 3: 7.65 ms per loop

Não use filtercomo um nome de variável, é uma função interna.


@ Mehdi Acho o caminho do Matlab altamente pouco intuitivo, mas acho que depende do que você está acostumado.
Ian Goldby

Como posso selecionar [2, 6]?
Florent

Eu entendo, list(compress(list_a, [not i for i in fill]))deve retornar[2, 6]
Florent

42

Igual a:

filtered_list = [i for (i, v) in zip(list_a, filter) if v]

O uso zipé a maneira pitônica de iterar várias seqüências em paralelo, sem a necessidade de indexação. Isso pressupõe que as duas seqüências tenham o mesmo comprimento (o zip pára após o menor tempo). Usar itertoolspara um caso tão simples é um pouco exagerado ...

Uma coisa que você faz no seu exemplo que realmente deveria parar de fazer é comparar as coisas com True, isso geralmente não é necessário. Em vez de if filter[idx]==True: ..., você pode simplesmente escrever if filter[idx]: ....


40

Com numpy:

In [128]: list_a = np.array([1, 2, 4, 6])
In [129]: filter = np.array([True, False, True, False])
In [130]: list_a[filter]

Out[130]: array([1, 4])

ou veja a resposta de Alex Szatmary se list_a puder ser um array numpy, mas não filtrar

Numpy geralmente também oferece um grande aumento de velocidade

In [133]: list_a = [1, 2, 4, 6]*10000
In [134]: fil = [True, False, True, False]*10000
In [135]: list_a_np = np.array(list_a)
In [136]: fil_np = np.array(fil)

In [139]: %timeit list(itertools.compress(list_a, fil))
1000 loops, best of 3: 625 us per loop

In [140]: %timeit list_a_np[fil_np]
10000 loops, best of 3: 173 us per loop

Bom ponto, eu prefiro usar o NumPymais listpossível. Mas se você precisar usar de listqualquer maneira, você (usando a NumPysolução) cria a np.arraypartir de ambas as listas, usa a indexação booleana e finalmente converte a matriz de volta à lista com o tolist()método Para ser mais preciso, você deve incluir a criação desses objetos na comparação de tempos. Então, o uso itertools.compressainda será a solução mais rápida.
Nerxis 9/07

17

Para fazer isso usando numpy, ou seja, se você tiver uma matriz a, em vez de list_a:

a = np.array([1, 2, 4, 6])
my_filter = np.array([True, False, True, False], dtype=bool)
a[my_filter]
> array([1, 4])

3
Se você transformar my_filter em uma matriz booleana, poderá usar a indexação booleana direta, sem a necessidade de where.
Bas Swinckels

1
filtered_list = [list_a[i] for i in range(len(list_a)) if filter[i]]

-1

Com o python 3, você pode usar list_a[filter]para obter Truevalores. Para obter Falsevalores, uselist_a[~filter]

Ao utilizar nosso site, você reconhece que leu e compreendeu nossa Política de Cookies e nossa Política de Privacidade.
Licensed under cc by-sa 3.0 with attribution required.