Splitting Numpy array based on value

Question

Suppose I have this NumPy array:

a = np.array([0, 3, 5, 5, 0, 10, 14, 15, 56, 0, 12, 23, 45, 23, 12, 45, 
              0, 1, 0, 2, 3, 4, 0, 0 ,0])

I would like to print all the numbers between 0s and automatically add them to a new np.array (see below):

a1=[3, 5, 5]
a2=[10, 14, 15, 56]
a3=[12, 23, 45, 23, 12, 45]
a4=[1]
a5=[2, 3, 4]

Is there a built-in function to do this?

Answer 1

Here's a vectorized approach using np.where and np.split -

idx = np.where(a!=0)[0]
aout = np.split(a[idx],np.where(np.diff(idx)!=1)[0]+1)

Sample run -

In [23]: a
Out[23]: 
array([ 0,  3,  5,  5,  0, 10, 14, 15, 56,  0,  0,  0, 12, 23, 45, 23, 12,
       45,  0,  1,  0,  2,  3,  4,  0,  0,  0])

In [24]: idx = np.where(a!=0)[0]

In [25]: np.split(a[idx],np.where(np.diff(idx)!=1)[0]+1)
Out[25]: 
[array([3, 5, 5]),
 array([10, 14, 15, 56]),
 array([12, 23, 45, 23, 12, 45]),
 array([1]),
 array([2, 3, 4])]

Answer 2

You can use groupby() function from itertools, and specify the key as the boolean condition of zero or nonzero. In such a way, all consecutive zeros and nonzeros will be grouped together. Use if filter to pick up groups of nonzeros and use list to convert the non zero groupers to lists.

from itertools import groupby
[list(g) for k, g in groupby(a, lambda x: x != 0) if k]

# [[3, 5], [10, 14, 15, 56], [12, 23, 45, 23, 12, 45], [1], [2, 3, 4]]