Answer by Ants Aasma for How can I partition (split up, divide) a list based on a condition?

Here's the lazy iterator approach:

from itertools import teedef 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)

It evaluates the condition once per item and returns two generators, first yielding values from the sequence where the condition is true, the other where it's false.

Because it's lazy you can use it on any iterator, even an infinite one:

from itertools import count, islicedef 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)))

Usually though the non-lazy list returning approach is better:

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

Edit: For your more specific usecase of splitting items into different lists by some key, heres a generic function that does that:

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

Usage:

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']