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

Turn a flat list into a nested list, with a specified number of lists per nesting level.

Python, 17 lines
 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[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]).

3 comments

Gabriel Genellina 14 years ago  # | flag

Instead of 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 range(12) example above.

Sander Evers (author) 14 years ago  # | flag

Thanks, I didn't know islice yet - coming from a functional programming background I scanned the itertools module for take but 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]) or 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 range(3)[:5] or islice(xrange(3),5) are ok).

Gabriel Genellina 14 years ago  # | flag

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)

Created by Sander Evers on Mon, 22 Feb 2010 (MIT)
Python recipes (4591)
Sander Evers's recipes (6)

Required Modules

Other Information and Tasks