Search code examples
pythonarrayscluster-computingpython-itertools

group clusters of numbers in array

I have an array like:

A = [1,3,8,9,3,7,2,1,3,9,6,8,3,8,8,1,2]

And I want to count the number of "entry clusters" that are >5. In this case the result should be 4, because:

[1, 3, (8,9), 3, (7), 2, 1, 3, (9,6,8), 3, (8,8), 1, 2]

Given L length of the array, I can do:

A = [1,3,8,9,3,7,2,1,3,9,6,8,3,8,8,1,2]
A = np.array(A)

for k in range(0,L):
    if A[k]>5:
        print k, A[k]

and this gives me all entries greater than 5. But how could I group every cluster of numbers?

Solution

  • You could use the groupby function from itertools.

    from itertools import groupby

A = [1,3,8,9,3,7,2,1,3,9,6,8,3,8,8,1,2]
result = [tuple(g) for k, g in groupby(A, lambda x: x > 5) if k]

print(result)
# [(8, 9), (7,), (9, 6, 8), (8, 8)]

print(len(result))
# 4