flatten nested lists

mawe 1 Tallied Votes 828 Views Share

The name says it all, I think. You have a nested list, you want to flatten it (e.g. because you want to count the occurence of an element), this function is all you need ;)

def flatten(lst):
    for elem in lst:
        if type(elem) in (tuple, list):
            for i in flatten(elem):
                yield i
        else:
            yield elem

# example usage
>>> nested = [('a', 'b', ['c', 'd'], ['a', 'b']), ['e', 'f']]
>>> nested.count("a")
 0   # woops!
>>> flattened = list( flatten(nested) )
>>> flattened
 ['a', 'b', 'c', 'd', 'a', 'b', 'e', 'f']

>>> flattened.count("a")
 2   # that's better

# with flatten() you can do everything you can do with generators,
# like such boring things
>>> for elem in flatten(nested):
....          print elem     
....
a
b
c
d
a
b
e
f