Retrieve intervals from array based on multiple ranges

user_11545694 0
Viewed 93

Let's say I have a Numpy array called a:

a = np.array([2,3,8,11,30,39,44,49,55,61])

I would like to retrieve multiple intervals based on two other arrays:

l = np.array([2,5,42])
r = np.array([10,40,70])

Doing something equivalent to this:

a[(a > l) & (a < r)]

With this as the desired output:

Out[1]: [[3 8],[ 8 11 30 39],[44 49 55 61]]

Of course I could do a simple for loop iterating over l and r, but the real life dataset is huge, so I would like to prevent looping as much as possible.

intervals numpy arrays python
2 Answers

You can't avoid looping given the ragged nature of output. But we should try to reduce compute when iterating. So, here's one way to simply slice into the input array while iterating, as we will most of the compute part with getting the start,stop indices per group with searchsorted -

lidx = np.searchsorted(a,l,'right')
ridx = np.searchsorted(a,r,'left')
out = [a[i:j] for (i,j) in zip(lidx,ridx)]
user_32938810

Here's one approach, broadcasting to obtain the indexing arrays, and using np.split to split the array:

# generates a (3,len(a)) where the windows are found in each column
w = (a[:,None] > l) & (a[:,None] < r)
# indices where in the (3,len(a)) array condition is satisfied
ix, _ = np.where(w)
# splits according to the sum along the columns
np.split(a[ix], np.cumsum(w.sum(0)))[:-1]
# [array([3, 8]), array([ 8, 11, 30, 39]), array([44, 49, 55, 61])]
user_96986840
Related