Turn a flat list into a nested list, with a specified number of lists per nesting level.
Python, 17 lines
Download
Copy to clipboard1 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[0]))
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[0]):
first,rest = _nest2(rest,levels[1:])
ret.append(first)
return ret,rest
else:
return flat[0],flat[1:]
Note: to return nested iterators instead of nested lists, simply replace list(itertools.islice(it,levels[0])) by itertools.islice(it,levels[0]).
Instead of
take(n,it)I'd probably uselist(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 therange(12)example above.Thanks, I didn't know
isliceyet - coming from a functional programming background I scanned the itertools module fortakebut 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(count(),[2,2,3])ornest(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 likerange(3)[:5]orislice(xrange(3),5)are ok).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)