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

An analogue of enumerate for nested lists. (Python recipe) by John Crichton
ActiveState Code (http://code.activestate.com/recipes/578376/)

A generator function which enables traversal and modification of deeply nested lists. Together with the supplied helper functions it could be useful when working with data stored in deeply nested lists, particularly when the level of nesting precludes a recursive approach.

Python, 80 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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
from functools import reduce
from collections import deque
from operator import getitem, setitem

def nested_enumerate(lst):
    """An analogue of enumerate for nested lists. 

       Returns an iterator over the (index, element) pairs of `lst` where 
       `index` is a list of integers [x0, x1,.., xn] such that 
       `lst[x0][x1]...[xn]==element`

       
       >>> for i, e in nested_enumerate([0, [[[1, [2, [[[3]]]]]]], [[[4]]]]):
               print('%s %s'%(str(i), str(e)))
       [0] 0
       [1, 0, 0, 0] 1
       [1, 0, 0, 1, 0] 2
       [1, 0, 0, 1, 1, 0, 0, 0] 3
       [2, 0, 0, 0] 4
    """
    
    # initial, partial index of lst
    partial_index = deque([([i], e) for (i, e) in enumerate(lst)])
      
    while partial_index:
        index, obj = partial_index.popleft()
        if isinstance(obj, list):
            # if obj is a list then its elements require further indexing
            new_dimension = [(index+[i], e) for (i, e) in enumerate(obj)]
            partial_index.extendleft(reversed(new_dimension)) 
        else:
            # obj is fully indexed
            yield index, obj


# complementary functions #

def nested_getitem(lst, index):
    """Returns lst[index[0]]...[index[n]]"""
    return reduce(getitem, index, lst)


def nested_setitem(lst, index, value):
    """Equivalent to the statement lst[index[0]]...[index[n]]=value"""
    setitem(
        reduce(getitem, index[0:-1], lst), index[-1], value
    )


# quick test #

deeplist = [0, [[[1, [2, [[[3]]]]]]], [[[4]]]]

for index, element in nested_enumerate(deeplist):
    assert nested_getitem(deeplist, index)==element

# example usage: applying a function to each element in a nested list #

square = lambda x: x**2

for index, element in nested_enumerate(deeplist):
    nested_setitem(deeplist, index, square(element))

assert deeplist==[0, [[[1, [4, [[[9]]]]]]], [[[16]]]]

# not recommended, but demonstrates different ways of traversing a list
# (plus, we all love flatten, right? ;-)

def flatten(lst):
    return [e for (i, e) in nested_enumerate(lst)]

def flatten2(lst):
    return [nested_getitem(lst, i) for (i, e) in nested_enumerate(lst)]

assert flatten(deeplist)==flatten2(deeplist)==[0, 1, 4, 9, 16]

# sort elements based on their depth of nesting, with deepest first
depthfirst = [e for (i, e) in sorted(nested_enumerate(deeplist), key=lambda (i, e):-len(i))]

assert depthfirst == [9, 4, 1, 16, 0]
Tags: enumerate, index, iterator, list