web-dev-qa-db-fra.com

Compréhensions des listes Python pour créer plusieurs listes

Je veux créer deux listes listOfA et listOfB pour stocker des index de A et B à partir d'une autre liste s.

s=['A','B','A','A','A','B','B']

La sortie devrait être deux listes

listOfA=[0,2,3,4]
listOfB=[1,5,6]

Je suis capable de faire cela avec deux déclarations.

listOfA=[idx for idx,x in enumerate(s) if x=='A']
listOfB=[idx for idx,x in enumerate(s) if x=='B']

Cependant, je veux le faire en une seule itération en utilisant uniquement des listes de compréhension. Est-il possible de le faire en une seule déclaration? Quelque chose comme listOfA,listOfB=[--code goes here--]

14
Heisenberg

La définition même d'une compréhension de liste est de produire un objet one list. Vos 2 objets de liste ont même des longueurs différentes; vous auriez à utiliser des effets secondaires pour réaliser ce que vous voulez.

N'utilisez pas de compréhension de liste ici. Il suffit d'utiliser une boucle ordinaire:

listOfA, listOfB = [], []

for idx, x in enumerate(s):
    target = listOfA if x == 'A' else listOfB
    target.append(idx)

Cela vous laisse juste avec one boucle à exécuter; cela dépassera toutes les compréhensions à deux listes, du moins pas jusqu'à ce que les développeurs trouvent un moyen de créer une liste deux fois plus rapide qu'une boucle avec des appels list.append() séparés.

Je choisirais ce jour-là une compréhension de liste imbriquée seulement pour pouvoir produire deux listes sur une ligne. Comme le Zen of Python l'indique:

La lisibilité compte.

36
Martijn Pieters

Sorte de; la clé est de générer une liste de 2 éléments que vous pouvez ensuite décompresser:

listOfA, listOfB = [[idx for idx, x in enumerate(s) if x == c] for c in 'AB']

Cela dit, je pense que c'est assez idiot de le faire de cette façon, une boucle explicite est beaucoup plus lisible.

10
RemcoGerlich

Une approche intéressante de ce problème consiste à utiliser defaultdict. Comme @Martin l'a déjà dit, la compréhension de liste n'est pas le bon outil pour produire deux listes. L'utilisation de defaultdict vous permettrait de créer une ségrégation en une seule itération. De plus, votre code ne serait limité sous aucune forme.

>>> from collections import defaultdict
>>> s=['A','B','A','A','A','B','B']
>>> listOf = defaultdict(list)
>>> for idx, elem in enumerate(s):
    listOf[elem].append(idx)
>>> listOf['A'], listOf['B']
([0, 2, 3, 4], [1, 5, 6])
5
Abhijit

Ce que vous essayez de faire n'est pas exactement impossible, c'est simplement compliqué et probablement inutile.

Si vous souhaitez partitionner un élément itérable en deux éléments, si la source est une liste ou un autre élément réutilisable, vous feriez probablement mieux de le faire en deux étapes, comme dans votre question. 

Même si la source est un itérateur, si la sortie souhaitée est une paire de listes, pas une paire d'itérateurs paresseux, utilisez la réponse de Martijn , ou faites deux passages sur list(iterator).)

Mais si vous avez vraiment besoin de partitionner un itératif arbitraire en deux itérables, il est impossible de le faire sans une sorte de stockage intermédiaire. 

Supposons que vous partitionnez [1, 2, -1, 3, 4, -2] en positives et negatives. Maintenant, vous essayez de next(negatives). Cela devrait vous donner -1, non? Mais cela ne peut pas être fait sans consommer le 1 et le 2. Ce qui signifie que lorsque vous essayez de next(positives), vous obtiendrez 3 au lieu de 1. Donc, le 1 et le 2 doivent être stockés quelque part.

La plupart de l'intelligence dont vous avez besoin se trouve à l'intérieur de itertools.tee . Si vous ne faites que positives et negatives en deux copies teed du même itérateur, filtrez-les tous les deux, vous avez terminé.

En fait, c’est l’une des recettes de la documentation itertools:

def partition(pred, iterable):
    'Use a predicate to partition entries into false entries and true entries'
    # partition(is_odd, range(10)) --> 0 2 4 6 8   and  1 3 5 7 9
    t1, t2 = tee(iterable)
    return filterfalse(pred, t1), filter(pred, t2)

(Si vous ne comprenez pas ça, ça vaut probablement la peine de l'écrire explicitement, avec deux fonctions génératrices partageant un itérateur et un tee via une fermeture, ou deux méthodes d'une classe partageant celles-ci via self. Cela devrait prendre deux douzaines de lignes. de code qui ne nécessite rien de délicat.)

Et vous pouvez même obtenir partition en tant qu'importation à partir d'une bibliothèque tierce telle que more_itertools.


Maintenant, vous pouvez utiliser ceci dans une ligne:

lst = [1, 2, -1, 3, 4, -2]
positives, negatives = partition(lst, lambda x: x>=0)

… Et vous avez un itérateur sur toutes les valeurs positives et un itérateur sur toutes les valeurs négatives. Ils semblent être complètement indépendants, mais ensemble, ils ne font qu'un seul passage sur lst - donc cela fonctionne même si vous affectez lst à une expression génératrice, à un fichier ou à quelque chose au lieu d'une liste.


Alors, pourquoi n'y a-t-il pas une sorte de syntaxe de raccourci pour cela? Parce que ce serait assez trompeur. 

Une compréhension ne prend pas de stockage supplémentaire. C'est pourquoi les expressions du générateur sont si géniales: elles peuvent transformer un itérateur paresseux en un autre itérateur paresseux sans rien stocker. 

Mais cela nécessite un stockage O(N). Imaginez que tous les chiffres soient positifs, mais vous essayez d’itérer negative en premier. Ce qui se produit? Tous les nombres sont poussés à trueq. En fait, cette O(N) pourrait même être infinite (par exemple, essayez-la sur itertools.count()).

Cela convient pour quelque chose comme itertools.tee, une fonction bloquée dans un module que la plupart des novices ne connaissent même pas, et qui a une belle documentation qui peut expliquer ce qu’elle fait et clarifier les coûts. Mais le faire avec du sucre syntaxique qui donnait l'impression que la compréhension normale serait une histoire différente.

1
abarnert

Pour ceux qui vivent au bord;)

listOfA, listOfB = [[i for i in cur_list if i is not None] for cur_list in Zip(*[(idx,None) if value == 'A' else (None,idx) for idx,value in enumerate(s)])]
0
Daniel Braun