Search code examples
pythonnumpyinsertion

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

Solution

  • 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