Splitting Numpy array based on value
NumPy's split()
and where()
in a list compehension:
[x[x!=0] for x in np.split(a, np.where(a==0)[0]) if len(x[x!=0])]
[array([3, 5, 5]),
array([10, 14, 15, 56]),
array([12, 23, 45, 23, 12, 45]),
array([1]),
array([2, 3, 4])]
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])]
You can get the indices of zeros with np.where:
zeros = np.where(a == 0)[0]
And iterate over every pair to slice the array:
[a[i+1:j] for i, j in zip(zeros, zeros[1:]) if len(a[i+1:j])>0]
Out[46]:
[array([3, 5]),
array([10, 14, 15, 56]),
array([12, 23, 45, 23, 12, 45]),
array([1]),
array([2, 3, 4])]
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]]