CODEWITHSUNDEEP
pythonnumpy
H.Bukhari
H.Bukhari

Reputation: 2291

Splitting Numpy array based on value

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?

Upvotes: 16

Views: 13052

Answers (6)

Sergey Zakharov
Sergey Zakharov

Reputation: 1615

This is probably the worst way to do this, but you could also convert your array into a string and then split that string a couple times:

long_string = "_".join(a.astype(str))

while long_string.startswith("0_"):
    long_string = long_string.removeprefix("0_")
while long_string.endswith("_0"):
    long_string = long_string.removesuffix("_0")

result = [list(map(int, i.split("_"))) for i in long_string.split("_0_")]

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

You would need Python 3.9 for .removeprefix() and .removesuffix().

Upvotes: 0

Divakar
Divakar

Reputation: 221684

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])]

Upvotes: 13

Mike Müller
Mike Müller

Reputation: 85582

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])]

Upvotes: 3

ox.
ox.

Reputation: 4005

No need for numpy, this lambda function works on a list, but we can convert your numpy array to and from a list on the way in and out:

cut = lambda x: [j for j in [cut(x[:x.index(0)])]+cut(x[x.index(0)+1:]) if j] if x.count(0) else x

numpy.array(cut(list(a)))

# array([[3, 5, 5], [10, 14, 15, 56], [12, 23, 45, 23, 12, 45], [1], [2, 3, 4]], dtype=object)

Upvotes: 0

akuiper
akuiper

Reputation: 215117

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]]

Upvotes: 7

user2285236
user2285236

Reputation:

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])]

Upvotes: 3

Related Questions