¿Cómo puedo particionar (dividir, dividir) una lista según una condición?
Tengo un código como:
good = [x for x in mylist if x in goodvals]
bad = [x for x in mylist if x not in goodvals]
El objetivo es dividir el contenido mylisten otras dos listas, en función de si cumplen o no una condición.
¿Cómo puedo hacer esto de manera más elegante? ¿ Puedo evitar hacer dos iteraciones separadas mylist? ¿Puedo mejorar el rendimiento al hacerlo?
Itere manualmente, utilizando la condición para seleccionar una lista a la que se agregará cada elemento:
good, bad = [], []
for x in mylist:
(bad, good)[x in goodvals].append(x)
good = [x for x in mylist if x in goodvals] bad = [x for x in mylist if x not in goodvals]¿Existe una forma más elegante de hacer esto?
¡Ese código es perfectamente legible y extremadamente claro!
# 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]
De nuevo, ¡esto está bien!
Puede haber ligeras mejoras de rendimiento al usar conjuntos, pero es una diferencia trivial, y encuentro que la comprensión de la lista es mucho más fácil de leer, y no tienes que preocuparte de que el orden se estropee, que se eliminen duplicados, etc.
De hecho, puedo dar un paso más "hacia atrás" y simplemente usar un bucle for simple:
images, anims = [], []
for f in files:
if f.lower() in IMAGE_TYPES:
images.append(f)
else:
anims.append(f)
La comprensión o el uso de una lista set()está bien hasta que necesite agregar alguna otra verificación u otro poco de lógica; digamos que desea eliminar todos los archivos JPEG de 0 bytes, simplemente agregue algo como...
if f[1] == 0:
continue
Aquí está el enfoque del iterador perezoso:
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)
Evalúa la condición una vez por elemento y devuelve dos generadores, primero generando valores de la secuencia donde la condición es verdadera y el otro donde es falsa.
Como es vago, puedes usarlo en cualquier iterador, incluso uno infinito:
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)))
Por lo general, aunque el enfoque de devolución de listas no diferidas es mejor:
def split_on_condition(seq, condition):
a, b = [], []
for item in seq:
(a if condition(item) else b).append(item)
return a, b
Editar: para su caso de uso más específico de dividir elementos en diferentes listas mediante alguna clave, aquí hay una función genérica que hace eso:
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
Uso:
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']
El problema con todas las soluciones propuestas es que escaneará y aplicará la función de filtrado dos veces. Haría una pequeña función simple como esta:
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 esa manera, no procesará nada dos veces y tampoco repetirá el código.
