ActiveState Code

Recipe 550805: Partitioning a sequence by a conditional function

This function takes a sequence and a function and returns two generators. The first generator yields all items in the sequence for which function(item) returns True, and the second generator yields all items for which function(item) returns False.

Python 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def partition(seq, func=bool):
    true_vals, false_vals = [], []
    def true_gen(_seq):
        while 1: 
            if true_vals:
                yield true_vals.pop(0)
            else:
                for item in _seq:
                    if func(item):
                        yield item
                        break
                    else:
                        false_vals.append(item)
                else:
                    break
    def false_gen(_seq):
        while 1:
            if false_vals:
                yield false_vals.pop(0)
            else:
                for item in _seq:
                    if not func(item):
                        yield item
                        break
                    else:
                        true_vals.append(item)
                else:
                    break
    it = iter(seq)
    return true_gen(it), false_gen(it)

Discussion

I created this function to emulate a function in Ruby that someone on #python needed.

Comments

  1. 1. At 3:52 a.m. on 2 mar 2008, Matteo Dell'Amico said:

    You can write it more simply by exploiting itertools.tee. Moreover, I would exploit the fact that False == 0 and True == 1, and extend it for n partitions.

    from itertools import tee

def partition(iterable, key=bool, n=2):
    decorated = ((key(item), item) for item in iterable)
    def filterkey(iterable, key):
        return (item for k, item in iterable if key == k)
    return [filterkey(it, i)
            for i, it in enumerate(tee(decorated, n))]

    By default, you get the same functionality as before, and you can write stuff like this:

    >>> partitions = partition(range(10), lambda x: x%3, 3)
>>> map(list, partitions)
[[0, 3, 6, 9], [1, 4, 7], [2, 5, 8]]
  2. 2. At 2:10 p.m. on 6 mar 2008, Raymond Hettinger said:

    Variant. def partition(seq, func=None): s1, s2 = tee(seq) return ifilter(func, s1), ifilterfalse(func, s2)

  3. 3. At 6:41 p.m. on 19 mar 2008, Aaron Gallagher (the author) said:

    The difference between my method and these other methods is that mine only iterates once over the sequence and doesn't require teeing the iterator. All-in-all, it probably uses less memory, though it might be longer.

  4. 4. At 4:57 a.m. on 8 may 2008, Paul Hankin said:

    Here's a general version of partition that doesn't tee the iterator.

    def partition(seq, func=bool, func_range=(True, False)):
    buffers = dict((x, []) for x in func_range)
    def values(i, seq=iter(seq)):
        while True:
            while not buffers[i]:
                x = seq.next()
                buffers[func(x)].append(x)
            yield buffers[i].pop(0)
    return tuple(values(i) for i in func_range)

print map(list, partition(range(10), lambda x: x % 3, range(3)))

Sign in to comment