find indices where one array is larger than an element in a second array

Question

I have two arrays

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

I want an output array with the length of b where each element b[i] is the index of the first element of a which is at least b[i]:

out = np.array([0, 5, 10, 15]

A slow solution is:

out = []
for x in b: 
    i = np.argmax( a >= x )
    out.append( i )

and this is a marginal speed increase:

out = []
i=0
for x in b: 
    i = np.argmax( a[i:] >= x ) + i
    out.append( i )

Any ideas for a pure numpy solution? This is prohibitively slow. Thanks

Bi Rico · Accepted Answer · 2019-01-23 04:32:57Z

3

If a is sorted, you can use a.searchsorted(b).

answered Jan 23, 2019 at 4:32

Bi Rico

25.9k3 gold badges57 silver badges75 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

thanks @Bi Rico! works great. Amazing the speed difference on 3e6 elements

In addition to being "pure numpy", searchsorted uses a binary search so it scales much better with the size of a. It's O(log(N)) instead of O(N).