Comment diviser une liste en fonction d'une condition ?

Question

Quelle est la meilleure façon, à la fois esthétiquement et du point de vue des performances, de diviser une liste d'éléments en plusieurs listes sur la base d'une condition ? L'équivalent de :

good = [x for x in mylist if x in goodvals]
bad  = [x for x in mylist if x not in goodvals]

Existe-t-il une manière plus élégante de procéder ?

Voici le cas d'utilisation réel, pour mieux expliquer ce que j'essaie de faire :

# files looks like: [ ('file1.jpg', 33L, '.jpg'), ('file2.avi', 999L, '.avi'), ... ]
IMAGE_TYPES = ('.jpg','.jpeg','.gif','.bmp','.png')
images = [f for f in files if f[2].lower() in IMAGE_TYPES]
anims  = [f for f in files if f[2].lower() not in IMAGE_TYPES]

Answer 1

good, bad = [], []
for x in mylist:
    (bad, good)[x in goodvals].append(x)

Answer 2

good = [x for x in mylist if x in goodvals]
bad  = [x for x in mylist if x not in goodvals]

Existe-t-il une manière plus élégante de procéder ?

Ce code est parfaitement lisible et extrêmement clair !

# files looks like: [ ('file1.jpg', 33L, '.jpg'), ('file2.avi', 999L, '.avi'), ... ]
IMAGE_TYPES = ('.jpg','.jpeg','.gif','.bmp','.png')
images = [f for f in files if f[2].lower() in IMAGE_TYPES]
anims  = [f for f in files if f[2].lower() not in IMAGE_TYPES]

Encore une fois, il s'agit très bien !

Il peut y avoir de légères améliorations de performance en utilisant des ensembles, mais c'est une différence triviale, et je trouve que la compréhension de la liste est beaucoup plus facile à lire, et vous n'avez pas à vous soucier de l'ordre, des doublons qui sont supprimés, etc.

En fait, je pourrais faire un pas de plus "en arrière" et utiliser une simple boucle "for" :

images, anims = [], []

for f in files:
    if f.lower() in IMAGE_TYPES:
        images.append(f)
    else:
        anims.append(f)

La compréhension d'une liste ou l'utilisation de set() est très bien jusqu'à ce que vous ayez besoin d'ajouter une autre vérification ou un autre élément de logique - disons que vous voulez supprimer tous les jpeg à 0 octet, il suffit d'ajouter quelque chose comme

if f[1] == 0:
    continue

Answer 3

Voici l'approche de l'itérateur paresseux :

from itertools import tee

def split_on_condition(seq, condition):
    l1, l2 = tee((condition(item), item) for item in seq)
    return (i for p, i in l1 if p), (i for p, i in l2 if not p)

Il évalue la condition une fois par élément et renvoie deux générateurs, le premier produisant des valeurs de la séquence lorsque la condition est vraie, l'autre lorsqu'elle est fausse.

Parce qu'il est paresseux, vous pouvez l'utiliser sur n'importe quel itérateur, même infini :

from itertools import count, islice

def is_prime(n):
    return n > 1 and all(n % i for i in xrange(2, n))

primes, not_primes = split_on_condition(count(), is_prime)
print("First 10 primes", list(islice(primes, 10)))
print("First 10 non-primes", list(islice(not_primes, 10)))

En général, l'approche consistant à renvoyer une liste non fastidieuse est meilleure :

def split_on_condition(seq, condition):
    a, b = [], []
    for item in seq:
        (a if condition(item) else b).append(item)
    return a, b

Edit : Pour votre cas d'utilisation plus spécifique qui consiste à diviser les éléments en différentes listes en fonction d'une clé, voici une fonction générique qui fait cela :

DROP_VALUE = lambda _:_
def split_by_key(seq, resultmapping, keyfunc, default=DROP_VALUE):
    """Split a sequence into lists based on a key function.

        seq - input sequence
        resultmapping - a dictionary that maps from target lists to keys that go to that list
        keyfunc - function to calculate the key of an input value
        default - the target where items that don't have a corresponding key go, by default they are dropped
    """
    result_lists = dict((key, []) for key in resultmapping)
    appenders = dict((key, result_lists[target].append) for target, keys in resultmapping.items() for key in keys)

    if default is not DROP_VALUE:
        result_lists.setdefault(default, [])
        default_action = result_lists[default].append
    else:
        default_action = DROP_VALUE

    for item in seq:
        appenders.get(keyfunc(item), default_action)(item)

    return result_lists

Utilisation :

def file_extension(f):
    return f[2].lower()

split_files = split_by_key(files, {'images': IMAGE_TYPES}, keyfunc=file_extension, default='anims')
print split_files['images']
print split_files['anims']

Answer 4

Le problème de toutes les solutions proposées est qu'elles analysent et appliquent la fonction de filtrage deux fois. Je ferais une petite fonction simple comme celle-ci :

def split_into_two_lists(lst, f):
    a = []
    b = []
    for elem in lst:
        if f(elem):
            a.append(elem)
        else:
            b.append(elem)
    return a, b

De cette manière, vous ne traitez rien deux fois et vous ne répétez pas le code.

Answer 5

J'aime beaucoup l'approche d'Anders car elle est très générale. Voici une version qui place le catégorisateur en premier (pour correspondre à la syntaxe du filtre) et qui utilise un dictionnaire par défaut (supposé importé).

def categorize(func, seq):
    """Return mapping from categories to lists
    of categorized items.
    """
    d = defaultdict(list)
    for item in seq:
        d[func(item)].append(item)
    return d