Search code examples
pythonfilterdata-partitioning

python equivalent of filter() getting two output lists (i.e. partition of a list)


Let's say I have a list, and a filtering function. Using something like

>>> filter(lambda x: x > 10, [1,4,12,7,42])
[12, 42]

I can get the elements matching the criterion. Is there a function I could use that would output two lists, one of elements matching, one of the remaining elements? I could call the filter() function twice, but that's kinda ugly :)

Edit: the order of elements should be conserved, and I may have identical elements multiple times.


Solution

  • Try this:

    def partition(pred, iterable):
        trues = []
        falses = []
        for item in iterable:
            if pred(item):
                trues.append(item)
            else:
                falses.append(item)
        return trues, falses
    

    Usage:

    >>> trues, falses = partition(lambda x: x > 10, [1,4,12,7,42])
    >>> trues
    [12, 42]
    >>> falses
    [1, 4, 7]
    

    There is also an implementation suggestion in itertools recipes:

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

    The recipe comes from the Python 3.x documentation. In Python 2.x filterfalse is called ifilterfalse.