Split NumPy array into individual parts with all zeros

Viewed 37

I have a one-dimensional np.array which has the following pattern

  1. array = [seq_1, seq_2, seq_3]
  2. seq_1 and seq_3 are all 0
  3. seq_2 contains at least a single 1

and I want to split array into seq_1, seq_2 and seq_3.

For example,

array = [0,0,0,1,1,1,1,0,0,0]
seq_1, seq_2, seq_3 = [0,0,0], [1,1,1,1], [0,0,0]
array = [0,1,0,0,0,0,0,0,1,0]
seq_1, seq_2, seq_3 = [0], [1,0,0,0,0,0,0,1], [0]
array = [0,1,1,1,0,1,0,1,1,0]
seq_1, seq_2, seq_3 = [0], [1,1,1,0,1,0,1,1], [0]

Is there a way using only NumPy functions (meaning no for loop) to achieve this?

1 Answers

You can find the indices to split on using np.diff and np.flatnonzero:

index = np.flatnonzero(np.diff(array)) + 1

Then you can split on the index:

seq = np.split(array, index)

If you want to unpack:

seq_1, seq_2, seq_3 = seq

As a one-liner:

seq_1, seq_2, seq_3 = np.split(array, np.flatnonzero(np.diff(array)) + 1)

This method will break up any array into homogeneous sequences, regardless of the values.

Related