7

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?

Vineet Jain
  • 1,515
  • 4
  • 21
  • 31
Eric So
  • 465
  • 2
  • 14

3 Answers3

11

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
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Do you perhaps have a suggestion how I could adapt your code such that I could do this row-wise on a 2D numpy array? – pr94 Mar 03 '21 at 15:48
5

Using groupby from itertools 

from itertools import groupby
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])
grouped_a = [ sum(1 for i in g) for k,g in groupby(a)]
Vineet Jain
  • 1,515
  • 4
  • 21
  • 31
0

I found a similar method to yours, just that this code finds the first and the last count separately. The answer is detailed in the code below:

import numpy as np

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])
print(f'a: {a}')

diff_a = np.diff(a)
print(f'diff_a: {diff_a}')

non_zero_pos_arr = np.where(diff_a != 0)[0]
print(f'Array of positions where non zero elements are present in diff_a array: {non_zero_pos_arr}')

diff_non_zero_pos_arr = np.diff(non_zero_pos_arr)
print(f'Result Array except for first and last element: {diff_non_zero_pos_arr}')

ans_first_ele = non_zero_pos_arr[0] + 1

ans_last_ele = len(diff_a) - non_zero_pos_arr[-1]

ans = np.array([], dtype=np.int8)
ans = np.append(ans, ans_first_ele)
ans = np.append(ans, diff_non_zero_pos_arr)
ans = np.append(ans, ans_last_ele)

print(f'Result Array: {ans}')

Output:

a: [0 1 1 1 0 0 0 0 0 0 0 1 0 1 1 0 0 0 1 1 0 0]
diff_a: [ 1  0  0 -1  0  0  0  0  0  0  1 -1  1  0 -1  0  0  1  0 -1  0]
Array of positions where non zero elements are present in diff_a array: 
 [ 0  3 10 11 12 14 17 19]
Result Array except for first and last element: [3 7 1 1 2 3 2]
Result Array: [1 3 7 1 1 2 3 2 2]
Rohan Katkam
  • 1