python - Getting the indices of several elements in a NumPy array at once

Question

Is there any way to get the indices of several elements in a NumPy array at once?

E.g.

import numpy as np
a = np.array([1, 2, 4])
b = np.array([1, 2, 3, 10, 4])

I would like to find the index of each element of a in b, namely: [0,1,4].

I find the solution I am using a bit verbose:

import numpy as np

a = np.array([1, 2, 4])
b = np.array([1, 2, 3, 10, 4])

c = np.zeros_like(a)
for i, aa in np.ndenumerate(a):
    c[i] = np.where(b==aa)[0]

print('c: {0}'.format(c))

Output:

c: [0 1 4]

Answer 1

You could use in1d and nonzero (or where for that matter):

>>> np.in1d(b, a).nonzero()[0]
array([0, 1, 4])

This works fine for your example arrays, but in general the array of returned indices does not honour the order of the values in a. This may be a problem depending on what you want to do next.

In that case, a much better answer is the one @Jaime gives here, using searchsorted:

>>> sorter = np.argsort(b)
>>> sorter[np.searchsorted(b, a, sorter=sorter)]
array([0, 1, 4])

This returns the indices for values as they appear in a. For instance:

a = np.array([1, 2, 4])
b = np.array([4, 2, 3, 1])

>>> sorter = np.argsort(b)
>>> sorter[np.searchsorted(b, a, sorter=sorter)]
array([3, 1, 0]) # the other method would return [0, 1, 3]