CODEWITHSUNDEEP
pythonnumpyinsertion
MrClean
MrClean

Reputation: 1460

Numpy get index of insertion in a sorted array

I am trying to find a way to create a function that passes two arrays, where the result is an array of the indices where the values from the first array will be located in the second array. The code below gives the result I want, but I am trying to get rid of the for loop and find a way to vectorize it using numpy functions:

x_array = np.array([25, 32, 3, 99, 300])
y_array = np.array([30, 33, 56, 99, 250])

result = [0, 1, 0, 3, -1]

def get_index(x_array, y_array):
   result = []
   for x  in x_array:
       index = np.where(x <= y_array)[0]
       if index.size != 0:
           result.append(index.min())
       else:
           result.append(-1)
   return result

Upvotes: 2

Views: 921

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114230

You are looking for np.searchsorted:

indices = np.searchsorted(y_array, x_array)

The only difference is that this returns the size of the array if you exceed the maximum element:

>>> indices
array([0, 1, 0, 3, 5], dtype=int64)

If you need to get -1 instead, you can use np.where or direct masking:

indices = np.where(indices < y_array.size, indices, -1)

OR

indices[indices >= y_array.size] = -1

Upvotes: 5

Related Questions