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]