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

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
 ``` 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] ```
 Created by John Crichton on Tue, 11 Dec 2012 (MIT)