Split numpy array into contiguous sections using numpy.where()-CodePudding

The function numpy.where() can be used to obtain an array of indices into a numpy array where a logical condition is true. How do we generate a list of arrays of the indices where each represents a contiguous region where the logical condition is true?

For example,

import numpy as np

a = np.array( [0.,.1,.2,.3,.4,.5,.4,.3,.2,.1,0.] )

idx = np.where( (np.abs(a-.2) <= .1) )

print( 'idx =', idx)
print( 'a[idx] =', a[idx] )

produces the following output,

idx = (array([1, 2, 3, 7, 8, 9]),)

a[idx] = [0.1 0.2 0.3 0.3 0.2 0.1]

and then

The question is, in a simple way, how do we obtain a list of arrays of indices, one such array for each contiguous section? For example, like this:

idx = (array([1, 2, 3]),), (array([7, 8, 9]),)

a[idx[0]] =  [0.1 0.2 0.3]

a[idx[1]] =  [0.3 0.2 0.1]

CodePudding user response：

This should work:

a = np.array([0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.4, 0.3, 0.2, 0.1, 0.0])
idx = np.nonzero(np.abs(a - 0.2) <= 0.1)[0]
splits = np.split(idx, np.nonzero(np.diff(idx) > 1)[0]   1)
print(splits)

It gives:

[array([1, 2, 3]), array([7, 8, 9])]

CodePudding user response：

You can check if the difference between the shifted idx arrays is 1, then split the array at the corresponding indices.

import numpy as np

a = np.array( [0.,.1,.2,.3,.4,.5,.4,.3,.2,.1,0.] )

idx = np.where( np.abs(a-.2) <= .1 )[0]

# Get the indices where the increment of values is larger than 1.
split_idcs = np.argwhere( idx[1:]-idx[:-1] > 1 )   1
# Split the array at the corresponding indices.
result = np.split(idx, split_idcs[0])

print(result)
# [array([1, 2, 3], dtype=int64), array([7, 8, 9], dtype=int64)]

It works for your example, however I am unsure if this implementation works for arbitrary sequences.

CodePudding user response：

You can achieved the goal by:

diff = np.diff(idx[0], prepend=idx[0][0]-1)
result = np.split(idx[0], np.where(diff != 1)[0])

idx = np.where((np.abs(a - .2) <= .1))[0]
diff = np.diff(idx, prepend=idx[0]-1)
result = np.split(idx, np.where(diff != 1)[0])

CodePudding user response：

You can simply use np.split() to split your idx[0] into contiguous runs:

ia = idx[0]
out = np.split(ia, np.where(ia[1:] != ia[:-1]   1)[0]   1)

>>> out
[array([1, 2, 3]), array([7, 8, 9])]