Search code examples
pythonperformancenumpyvectorization

vectorise for loop with numpy.where

>>> a = np.array([1, 1, 2, 3, 3, 4, 5])
>>> b = np.array([1, 3, 4, 5])
>>> indices = np.zeros((len(a)))
>>> for i in range(len(a)):
    try:
        indices[i] = np.where(b == a[i])[0][0]
    except:
        indices[i] = -1

>>> indices
array([ 0.,  0., -1.,  1.,  1.,  2.,  3.])

For each element of an np.array, how can I get its index in another np.array and also put a constant value (here, -1) where the element is not found in the latter one?

Solution

  • Use np.searchsorted to get the indices and then re-use those to get the invalid ones by comparing the indexed ones against a and set those as -1 -

    idx = np.searchsorted(b,a)
idx[b[idx] != a] =-1