Python: selecione o subconjunto da lista com base no conjunto de índices

98

Tenho várias listas com o mesmo número de entradas (cada uma especificando uma propriedade de objeto):

property_a = [545., 656., 5.4, 33.]
property_b = [ 1.2,  1.3, 2.3, 0.3]
...

e lista com sinalizadores do mesmo comprimento

good_objects = [True, False, False, True]

(que poderia ser facilmente substituído por uma lista de índice equivalente:

good_indices = [0, 3]

Qual é a maneira mais fácil de gerar novas listas property_asel, property_bsel... que contêm apenas os valores indicados quer pelas Trueentradas ou os índices?

property_asel = [545., 33.]
property_bsel = [ 1.2, 0.3]
Fuenfundachtzig
fonte

Respostas:

126

Você pode simplesmente usar a compreensão de lista :

property_asel = [val for is_good, val in zip(good_objects, property_a) if is_good]

ou

property_asel = [property_a[i] for i in good_indices]

O último é mais rápido porque há menos do good_indicesque o comprimento de property_a, supondo que good_indicessejam pré-computados em vez de gerados instantaneamente.


Editar : a primeira opção é equivalente ao itertools.compressdisponível desde Python 2.7 / 3.1. Veja a resposta de @Gary Kerr .

property_asel = list(itertools.compress(property_a, good_objects))
Kennytm
fonte
1
@fuen: Sim. Causa muito no Python 2 (em vez disso, use itertools.izip ), não tanto no Python 3. Isso ocorre porque o zipno Python 2 criará uma nova lista, mas no Python 3 retornará apenas um gerador (lento).
kennytm
OK, então devo me ater à sua segunda proposta, porque ela constitui a parte central do meu código.
Fuenfundachtzig
4
@ 85: por que você está se preocupando com o desempenho? Escreva o que você precisa fazer, se for lento, e teste para encontrar gargalos.
Gary Kerr
1
@PreludeAndFugue: Se houver duas opções equivalentes, é bom saber qual é a mais rápida e usá-la imediatamente.
fuenfundachtzig
1
Você pode apenas usar from itertools import izipe usar isso em vez do zipprimeiro exemplo. Isso cria um iterador, igual ao Python 3.
Chris B.
28

Vejo 2 opções.

  1. Usando numpy:

    property_a = numpy.array([545., 656., 5.4, 33.])
    property_b = numpy.array([ 1.2,  1.3, 2.3, 0.3])
    good_objects = [True, False, False, True]
    good_indices = [0, 3]
    property_asel = property_a[good_objects]
    property_bsel = property_b[good_indices]
  2. Usando uma lista de compreensão e compactá-la:

    property_a = [545., 656., 5.4, 33.]
    property_b = [ 1.2,  1.3, 2.3, 0.3]
    good_objects = [True, False, False, True]
    good_indices = [0, 3]
    property_asel = [x for x, y in zip(property_a, good_objects) if y]
    property_bsel = [property_b[i] for i in good_indices]
Wolph
fonte
2
Usar o Numpy é uma boa sugestão, pois o OP parece querer armazenar números em listas. Uma matriz bidimensional seria ainda melhor.
Philipp
Também é uma boa sugestão porque esta será uma sintaxe muito familiar aos usuários do R, onde este tipo de seleção é muito poderoso, principalmente quando aninhado e / ou multidimensional.
Thomas Browne
[property_b[i] for i in good_indices]é um bom para usar semnumpy
Ilya Rusin
16

Use a função integrada zip

property_asel = [a for (a, truth) in zip(property_a, good_objects) if truth]

EDITAR

Basta olhar para os novos recursos do 2.7. Agora existe uma função no módulo itertools que é semelhante ao código acima.

http://docs.python.org/library/itertools.html#itertools.compress

itertools.compress('ABCDEF', [1,0,1,0,1,1]) =>
  A, C, E, F
Gary Kerr
fonte
1
Não estou impressionado com o uso itertools.compressdaqui. A compreensão da lista é muito mais legível, sem ter que desenterrar o que essa compressão está fazendo.
PaulMcG
5
Hm, acho o código usando compress muito mais legível :) Talvez eu seja tendencioso, porque faz exatamente o que eu quero.
fuenfundachtzig
8

Supondo que você tenha apenas a lista de itens e uma lista de índices verdadeiros / obrigatórios, este deve ser o mais rápido:

property_asel = [ property_a[index] for index in good_indices ]

Isso significa que a seleção da propriedade só fará tantas rodadas quantos forem os índices verdadeiros / obrigatórios. Se você tiver muitas listas de propriedades que seguem as regras de uma única lista de tags (verdadeiro / falso), pode criar uma lista de índices usando os mesmos princípios de compreensão de lista:

good_indices = [ index for index, item in enumerate(good_objects) if item ]

Isso itera por meio de cada item em good_objects (enquanto lembra seu índice com enumerar) e retorna apenas os índices onde o item é verdadeiro.


Para quem não está entendendo a lista, aqui está uma versão em prosa em inglês com o código destacado em negrito:

liste o índice para cada grupo de índice, item que existe em uma enumeração de objetos bons , se (onde) o item é Verdadeiro

Eyrofire
fonte
0

As linguagens Matlab e Scilab oferecem uma sintaxe mais simples e elegante do que Python para a pergunta que você está fazendo, então acho que o melhor que você pode fazer é imitar Matlab / Scilab usando o pacote Numpy em Python. Ao fazer isso, a solução para o seu problema é muito concisa e elegante:

from numpy import *
property_a = array([545., 656., 5.4, 33.])
property_b = array([ 1.2,  1.3, 2.3, 0.3])
good_objects = [True, False, False, True]
good_indices = [0, 3]
property_asel = property_a[good_objects]
property_bsel = property_b[good_indices]

O Numpy tenta imitar o Matlab / Scilab, mas tem um custo: você precisa declarar todas as listas com a palavra-chave "array", algo que sobrecarregará seu script (esse problema não existe no Matlab / Scilab). Observe que esta solução é restrita a matrizes de número, que é o caso em seu exemplo.

FredAndre
fonte
3
Em nenhum lugar da pergunta ele menciona o NumPy - não há necessidade de expressar sua opinião sobre o NumPy vs Matlab. Listas Python não são a mesma coisa que matrizes NumPy, mesmo que ambas correspondam aproximadamente a vetores. (As listas Python são como matrizes de células Matlab - cada elemento pode ter um tipo de dados diferente. As matrizes NumPy são mais restritas para permitir certas otimizações). Você pode obter uma sintaxe semelhante ao seu exemplo por meio da filterbiblioteca interna ou externa do Python pandas. Se você vai trocar de idioma, também pode tentar R, mas não é isso que a pergunta está perguntando .
Livius