Question

Hi I have a python problem whereby I have to count which list elements contain the value 2, in a list with various levels of nest. E.g.:

my_list = [[2,3,2,2], [[2,1,2,1], [2,1,1]], [1,1,1]]

this list can have have up to three levels of nesting but could also be two or only one level deep.

I have piece of code which sort of works:

count = 0
for i in my_list:
    if len(i) > 1:
        for x in i:
            if 2 in x:
                count += 1
    elif i == 2:
        count += 1

However, apart from being very ugly, this does not take account of the potential to have a list with a single element which is 2. It also doesn’t work to get len() on a single int.

I know list comprehension should be able to take care of this but I am a litle stuck on how to deal with the potential nesting.

Any help would be much appreciated.

Answer 1

I’d use a variant of the flatten() generator from https://stackoverflow.com/a/2158532/367273

The original yields every elements from an arbitrarily nested and irregularly shaped structure of iterables. My variant (below) yields the innermost iterables instead of yielding the scalars.

from collections import Iterable

def flatten(l):
    for el in l:
        if isinstance(el, Iterable) and any(isinstance(subel, Iterable) for subel in el):
            for sub in flatten(el):
                yield sub
        else:
            yield el

my_list = [[2,3,2,2], [[2,1,2,1], [2,1,1]], [1,1,1]]
print(sum(1 for el in flatten(my_list) if 2 in el))

For your example it prints 3.