Welcome, guest | Sign In | My Account | Store | Cart

Partitioning a sequence by a conditional function (Python recipe) by Aaron Gallagher
ActiveState Code (http://code.activestate.com/recipes/550805/)

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, 30 lines
Download
 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)

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

Tags: algorithms

4 comments

Matteo Dell'Amico 16 years, 1 month ago  # | flag

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]]
Raymond Hettinger 16 years, 1 month ago  # | flag

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

Aaron Gallagher (author) 16 years, 1 month ago  # | flag

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.

Paul Hankin 15 years, 11 months ago  # | flag

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)))
Created by Aaron Gallagher on Sat, 1 Mar 2008 (PSF)
◄ Python recipes (4591) ►
◄ Aaron Gallagher's recipes (3) ►

Required Modules

  • (none specified)

Other Information and Tasks