Python iterate through list and count items of a certain value

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.

Answer 2

Update and final answer:

My solution recurses a nested list of lists. If the list contains 2, the count is set to one. Then the sum of the counts of any sub-lists are added. This approach supports heterogeneous lists where numbers and lists can be mixed at the same level, as in the second usage example below:

import collections

def list_count(l):
    count = int(2 in l)
    for el in l:
        if isinstance(el, collections.Iterable):
            count += list_count(el)
    return count

Here are a couple of test cases:

my_list = [[2,3,2,2], [[2,1,2,1], [2,1,1]], [1,1,1]]
print = list_count(my_list)
# 3 is printed

my_list = [2, [2,3,2,2], [[2,1,2,1], [2,1,1]], [1,1,1]]
print = list_count(my_list)
# 4 is printed

@Akavall's answer reminded me this could be collapsed to a one-liner. (But I always get (usually justified) readability complaints on SO when I do this.)

def list_count(l):
    return ( int(2 in l) +
        sum([list_count(el) for el in l 
            if isinstance(el, collections.Iterable)]) )

Original Answer (not what original question was looking for)

Update: @NPE updated his answer after expected result was specified. He's current answer works as the original poster desires.

@NPE's (original) answer is close, but you asked:

count which list elements contain the value 2

Which I read as you needing:

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

But he's looking for 3. My code generates 2 because it iterates across each top level element and counts it if it contains any 2, but that's not what he actually wants.

Answer 3

Here is another way to do this:

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

from collections import Iterable 

def count_lists_w_two(x):
    if isinstance(x, Iterable) == False:
        return 0
    else:
        return (2 in x) + sum(count_lists_w_two(ele) for ele in x)

Result:

>>> count_lists_w_two(small_list)
1
>>> count_lists_w_two(my_list)
3
>>> count_lists_w_two(another_list)
4