Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python - Finding each occurrence of a value in a mixed array (integers, lists)

I have an array: x = [ [1, 2], 1, 1, [2, 1, [1, 2]] ] in which I want to count every occurrence of the number 1, and store that number in the variable one_counter. x.count(1) returns only 2 occurrences of 1, which is insufficient.

My code below serves my purpose and stores 5 in one_counter, however it looks messy and feels unpythonic to me.

Any suggestions how I can improve its pythonicity and expand it into more-dimensional lists?

Thanks!

x = [[1, 2], 1, 1, [2, 1, [1, 2]]]

one_counter = 0

for i in x:
    if type(i) == list:
        for j in i:
            if type(j) == list:
                for k in j:
                    if k == 1:
                        one_counter += 1

            else:
                if j == 1:
                    one_counter += 1

    else:
        if i == 1:
            one_counter += 1
like image 351
Shon Freelen Avatar asked Dec 13 '22 08:12

Shon Freelen


1 Answers

You could use recursion:

def flatten_count(iterable, element):
    count = 0
    for item in iterable:
        if item == element:
            count += 1
        if isinstance(item, list):
            count += flatten_count(item, element)
    return count

Or more concisely:

def flatten_count(iterable, element):
    return sum(
        flatten_count(item, element) if isinstance(item, list) else item == element
        for item in iterable 
    )

Use like this:

>>> x = [[1, 2], 1, 1, [2, 1, [1, 2]]]
>>> print(flatten_count(x, 1))
5
like image 100
Mark Byers Avatar answered Feb 22 '23 23:02

Mark Byers