Ekeyme Mo Ekeyme Mo - 4 months ago 24
Python Question

Count elements in a nested list in an elegant way

I have nested tuples in a list like

l = [(1, 'a', 'b'), (2, 'b', 'c'), (3, 'e', 'a')]


I want to know how many 'a' and 'b' in the list in total. So I currently use the following code to get the result.

amount_a_and_b = len([None for _, elem2, elem3 in l if elem2 == 'a' or elem3 == 'b'])


But I got
amount_a_and_b = 1
, so how to get the right answer?

Also, is there a more elegant way (less code or higher performance or using builtins) to do this?

Answer

I'd flatten the list with itertools.chain.from_iterable() and pass it to a collections.Counter() object:

from collections import Counter
from itertools import chain

counts = Counter(chain.from_iterable(l))
amount_a_and_b = counts['a'] + counts['b']

Or use sum() to count how many times a value appears in the flattened sequence:

from itertools import chain

amount_a_and_b = sum(1 for v in chain.from_iterable(l) if v in {'a', 'b'})

The two approaches are pretty much comparable in speed on Python 3.5.1 on my Macbook Pro (OS X 10.11):

>>> from timeit import timeit
>>> from collections import Counter
>>> from itertools import chain
>>> l = [(1, 'a', 'b'), (2, 'b', 'c'), (3, 'e', 'a')] * 1000  # make it interesting
>>> def counter():
...     counts = Counter(chain.from_iterable(l))
...     counts['a'] + counts['b']
...
>>> def summing():
...     sum(1 for v in chain.from_iterable(l) if v in {'a', 'b'})
...
>>> timeit(counter, number=1000)
0.5640139860006457
>>> timeit(summing, number=1000)
0.6066895100011607