Count consecutive occurrences of values ​​varying in length in a numpy array

Question

Let's say I have a set of numbers in a numpy array, and I'm testing them based on a condition that returns a boolean array:

np.random.seed(3456) a = np.random.rand(8) condition = a>0.5 

And with this boolean array, I want to calculate all the lengths of consecutive True inputs. For example, if I had [True,True,True,False,False,True,True,False,True] , I would like to return [3,2,1] .

I can do this with this code:

 length,count = [],0 for i in range(len(condition)): if condition[i]==True: count += 1 elif condition[i]==False and count>0: length.append(count) count = 0 if i==len(condition)-1 and count>0: length.append(count) print length 

But is there anything already implemented for this function or a python, numpy, scipy, etc. function that counts the length of consecutive entries in a list or array for a given input?

Answer 1

Here's a solution using itertools (this is probably not the fastest solution):

 import itertools condition = [True,True,True,False,False,True,True,False,True] [ sum( 1 for _ in group ) for key, group in itertools.groupby( condition ) if key ] Out: [3, 2, 1] 

Answer 2

If you already have a numpy array, this will probably be faster:

 >>> condition = np.array([True,True,True,False,False,True,True,False,True]) >>> np.diff(np.where(np.concatenate(([condition[0]], condition[:-1] != condition[1:], [True])))[0])[::2] array([3, 2, 1]) 

It discovers where the fragments start, has some logic for the first and last fragments, and simply calculates the differences between starts and discards the lengths corresponding to False pieces.