Turn a flat list into a nested list, with a specified number of lists per nesting level.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
# Example: # >>> nest(range(12),[2,2,3]) # [[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]] def nest(flat,levels): '''Turn a flat list into a nested list, with a specified number of lists per nesting level. Excess elements are silently ignored.''' return _nest(flat,levels).next() def _nest(flat,levels): if levels: it = _nest(flat,levels[1:]) while 1: yield list(itertools.islice(it,levels)) else: for d in flat: yield d
The code uses generators, and is faster than if function calls were used; in the [2,2,3] example, _nest is only called 3 times (once per nesting level) instead of 2x2x3 times --- compare against:
def nest2(flat,levels): first,rest = _nest2(flat,levels) return first def _nest2(flat,levels): if levels: ret =  rest = flat for i in xrange(levels): first,rest = _nest2(rest,levels[1:]) ret.append(first) return ret,rest else: return flat,flat[1:]
Note: to return nested iterators instead of nested lists, simply replace
take(n,it)I'd probably use
list(itertools.islice(it,n))to avoid building those intermediate lists.
Also, one should warn potential users that excess elements are discarded:
nest(range(100),[2,2,3])gives the very same result as the
Thanks, I didn't know
isliceyet - coming from a functional programming background I scanned the itertools module for
takebut didn't find it, and defined it myself. I've updated the recipe accordingly.
The silent ignoring of excess elements can also be used as a feature, for example in
nest(repeat(0),[2,2,3]), so I'm in doubt whether it should raise an exception. Also, trying to nest a flat list with too few elements is allowed (just like
Yes, itertools terminology is slightly different than, say, Haskell. It has takewhile and dropwhile, but no take nor drop... islice is like a generalization of take (sort of).
The way excess elements (or missing elements) are handled is perfectly fine with me - I didn't meant to raise an exception or anything, just asked to clarify it (it isn't obvious what happens in such cases, just by looking at the code)