Wednesday, October 24, 2012

Splitting a python list into groups of n

Say you have a list L, that you want broken into sublists l1, l2, ... lm of length n. What's the best way to go about doing this in python?

One simple idiom works as follows:

>>> groupn = lambda l, n: zip(*(iter(l),) * n)
>>> groupn(range(12), 3)
[(0, 1, 2), (3, 4, 5), (6, 7, 8), (9, 10, 11)]

So, how exactly does this work?

First, let's take a look at an alternate, but more syntactically revealing version of the groupn function:
>>> alt_groupn = lambda l, n: zip(*[iter(l)] * n)
>>> alt_groupn(range(12), 3)
[(0, 1, 2), (3, 4, 5), (6, 7, 8), (9, 10, 11)]

The key thing to note here is the [iter(l)] * n. This is creating a list of length n of references to the *same* iterator.
>>> l = range(10)
>>> [ iter(l) ] * 3
[<listiterator object at 0x10a1d6310>, <listiterator object at 0x10a1d6310>, <listiterator object at 0x10a1d6310>]
Unpacking this list as an argument to zip is equivalent to:
>>> i = iter(range(10))
>>> zip(i, i, i)
[(0, 1, 2), (3, 4, 5), (6, 7, 8)]
The magic happens because zip draws elements from the same iterator to fill the different positions of each tuple!

No comments:

Post a Comment