Is there a way to "fork" a list in two based on a condition

Question

I have seen a pattern repeated a couple times in my team's code, it looks like this

numbers = [1, 2, 3, 4]
even_numbers = [n for n in numbers if n % 2 == 0]
odd_numbers = [n for n in numbers if n % 2 != 0]

I was wondering if there is a function somewhere (I have looked around but haven't been able to find it) that would do something like this

numbers = [1, 2, 3, 4]
even_numbers, odd_numbers = fork(numbers, lambda x: x % 2 == 0)

So, this function I am looking for, would receive an iterable and a function, and return two lists, one would be the values that match a provided condition, and the other would be the ones that didn't.

Is there something around python's standard library that achieves this?

Normally this is called a "partition", there's a [recipe in `itertools`](https://docs.python.org/3/library/itertools.html#itertools-recipes) for it. — jonrsharpe, Aug 08 '18 at 09:09

AKX · Accepted Answer · 2018-08-08T10:04:23.767

11

I usually call this sift, but partition is fine too.

Another, itertools-less implementation might be

def sift(iterable, predicate):
    t = []
    f = []
    for value in iterable:
        (t if predicate(value) else f).append(value)
    return (t, f)

even, odd = sift([1, 2, 3, 4, 5], lambda x: x % 2 == 0)

EDIT: for a slightly more complex implementation that is about 30% faster (on my Python installation anyway):

def sift2(iterable, predicate):
    t = []
    f = []
    ta = t.append
    fa = f.append
    for value in iterable:
        (ta if predicate(value) else fa)(value)
    return (t, f)

edited Aug 08 '18 at 10:04

answered Aug 08 '18 at 09:23

AKX

152,115
15
115
172

I really like this implementation, I will mark it as the answer since it does not use itertools and iterates only once. – Jose Enrique Narvaez Aug 08 '18 at 09:25
using `timeit` suggests this takes quite a bit longer than the original approach for a list of 1000 items. `%timeit two_lists(lst) 186 µs ± 526 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)` vs `%timeit using_sift(lst) 290 µs ± 360 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)` though it does have an additional function call – roganjosh Aug 08 '18 at 09:56
@roganjosh Yeah, that's probably the function call overhead. Can't do much about it for a generic solution. – AKX Aug 08 '18 at 10:00
@AKX I'll modify the code a bit and time without function call over head. Well, if I can find an elegant way to balance the books a bit :) – roganjosh Aug 08 '18 at 10:01

score 3 · Answer 2 · answered Aug 08 '18 at 09:22

You can use the following function:

from itertools import filterfalse, tee

def fork(pred, iterable):
    'Use a predicate to partition entries into false entries and true entries'
    t1, t2 = tee(iterable)
    return list(filterfalse(pred, t1)), list(filter(pred, t2))

Source: itertools

user2653663 · Answer 3 · 2018-08-08T09:32:09.250

Full code following @jonrsharpe suggestion.

import itertools

def fork(iterable):
    "Returns list of even, odd elements of a list"
    t1, t2 = itertools.tee(iterable)
    pred = lambda i: i % 2 == 0
    return list(filter(pred, t2)), list(itertools.filterfalse(pred, t1))

odd, even = fork([1,2,3,4,5])
print(odd)
print(even)

Alternative numpy version which might be faster for big arrays

import numpy as np

def fork(iterable):
    "Returns array of even, odd elements of an array"
    iterable_array = np.asarray(iterable)
    mask = (iterable_array % 2 == 0)
    return iterable_array[~mask], iterable_array[mask]

score 2 · Answer 4 · answered Aug 08 '18 at 09:27

I did not find anything in standard library performing what you want. I suggest you this user-defined implementation, which is not optimized at all but very simple and easy to read:

def myFunc(iterable, func):

    first  = [i for i in iterable if func(i)]
    second = [i for i in iterable if not func(i)]

    return first,second


numbers = [1, 2, 3, 4]
even_numbers, odd_numbers = myFunc(numbers, lambda x: x % 2 == 0)

print(even_numbers) # [2, 4]
print(odd_numbers)  # [1, 3]

score 0 · Answer 5 · answered Aug 08 '18 at 09:23

You can create your own function:

l = [1, 2, 3, 4]
def fork(l,key):
   return list(filter(key,l)), [i for i in l if i not in list(filter(key,l))]
even_numbers, odd_numbers = fork(l, lambda x: x % 2 == 0)
print(even_numbers)
print(odd_numbers)

Output:

[2, 4]
[1, 3]

Is there a way to "fork" a list in two based on a condition

5 Answers5