[FIXED] Split numpy array according to array containing bounds without for loop

Issue

Let's say I have an 1-D array with some values as

a = np.array([1,2,3,5,6,8,9,11,12,13,14,15])

I also have another array with some values as

b = np.array([0,3,6,9,12,15,18])

I need to split array a in such a way that the first array should have elements with values between b[0] and b[1], the second array should have elements with values between b[1] and b[2], and so on. So the results would be the following

result = [[1,2],
          [3,5],
          [6,8],
          [9,11],
          [12,13,14],
          [15]]

I know how to do this using a for loop but that's way too slow. Would there be a faster way?

Solution

Try as follows. First, we use np.searchsorted with side parameter set to left to "[f]ind indices where elements [from a] should be inserted to maintain order [inside b]". So:

import numpy as np

a = np.array([1,2,3,5,6,8,9,11,12,13,14,15])
b = np.array([0,3,6,9,12,15,18])

a_in_b = np.searchsorted(a, b, side='left')
# array([ 0,  2,  4,  6,  8, 11, 12], dtype=int64)

Next, we feed this array to np.split as the indices_or_sections parameter, and receive a list of np.ndarrays. Finally, we cut off the first and last ones, as they will be empty (i.e. array([], dtype=int32)).

res = np.split(a, a_in_b, axis=0)[1:-1]
print(res)

[array([1, 2]),
 array([3, 5]),
 array([6, 8]),
 array([ 9, 11]),
 array([12, 13, 14]),
 array([15])]

Let me know if you run into any complications applying this method to your actual data.

Answered By - ouroboros1