python - How to count continuous numbers in numpy

Question

I have a Numpy one-dimensional array of 1 and 0. for e.g

a = np.array([0,1,1,1,0,0,0,0,0,0,0,1,0,1,1,0,0,0,1,1,0,0])

I want to count the continuous 0s and 1s in the array and output something like this

[1,3,7,1,1,2,3,2,2]

What I do atm is

np.diff(np.where(np.abs(np.diff(a)) == 1)[0])

and it outputs

array([3, 7, 1, 1, 2, 3, 2])

as you can see it is missing the first count 1.

I've tried np.split and then get the sizes of each segments but it does not seem to be optimistic.

Is there more elegant "pythonic" solution?

Answer 1

Here's one vectorized approach -

np.diff(np.r_[0,np.flatnonzero(np.diff(a))+1,a.size])

Sample run -

In [208]: a = np.array([0,1,1,1,0,0,0,0,0,0,0,1,0,1,1,0,0,0,1,1,0,0])

In [209]: np.diff(np.r_[0,np.flatnonzero(np.diff(a))+1,a.size])
Out[209]: array([1, 3, 7, 1, 1, 2, 3, 2, 2])

Faster one with boolean concatenation -

np.diff(np.flatnonzero(np.concatenate(([True], a[1:]!= a[:-1], [True] ))))

Runtime test

For the setup, let's create a bigger dataset with islands of 0s and 1s and for a fair benchmarking as with the given sample, let's have the island lengths vary between 1 and 7 -

In [257]: n = 100000 # thus would create 100000 pair of islands

In [258]: a = np.repeat(np.arange(n)%2, np.random.randint(1,7,(n)))

# Approach #1 proposed in this post
In [259]: %timeit np.diff(np.r_[0,np.flatnonzero(np.diff(a))+1,a.size])
100 loops, best of 3: 2.13 ms per loop

# Approach #2 proposed in this post
In [260]: %timeit np.diff(np.flatnonzero(np.concatenate(([True], a[1:]!= a[:-1], [True] ))))
1000 loops, best of 3: 1.21 ms per loop

# @Vineet Jain's soln    
In [261]: %timeit [ sum(1 for i in g) for k,g in groupby(a)]
10 loops, best of 3: 61.3 ms per loop