Home > Enterprise >  Separating elements from a list in Python depending on a condition
Separating elements from a list in Python depending on a condition

Time:01-23

I have a list of elements and want to separate the elements of the list by a certain condition.

A simple example is a list of numbers and i want to separate the odd from the even ones. For that could use the filter builtin like so:

def is_even(x):
    # ...

l = [0, 1, 2, 3, 4, 5, 6]

even = list(filter(is_even, l))
odd = list(filter(not is_even, l))

That is a bit error prone if the condition is a bit more complex, because i repeat myself twice in the filter functions. Is there a more elegant way to achieve this?

CodePudding user response:

itertools has a recipe exactly for that:

from itertools import tee, filterfalse
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)

Usage:

odd, even = partition(is_even, l)

You can convert them to lists, but I suggest leaving them as iterators.

CodePudding user response:

If you do not wish to run the predicate twice, you can create 2 lists like so:

def split_predicate(pred, iterable):
    """Split an iterable into two lists based on a predicate.
    
    The predicate will only be called once per element.
    
    Returns:
        A tuple of two lists, the first containing all elements for which
        the predicate returned True, the second containing all elements for
        which the predicate returned False."""
    t1, t2 = [], []
    for item in iterable:
        (t1 if pred(item) else t2).append(item)
    return t1, t2

You can also optimize it using generators. Tell me if you wish that I'll write it, it's a bit more complex.

  • Related