CODEWITHSUNDEEP
numpy
Fenchurch
Fenchurch

Reputation: 303

Find indices of a list of values in a numpy array

I have a numpy master array. Given another array of search values, with repeating elements, I want to produce the indices of these search values in the master array.

E.g.: master array is [1,2,3,4,5], search array is [4,2,2,3]

Solution: [3,1,1,2]

Is there a "native" numpy function that does this efficiently (meaning at C speed, rather than python speed)?

I'm aware of the following solution, but, first, it's a python list comprehension, and second, it'll search for the index of 2 twice.

ma = np.array([1,2,3,4,5])
sl = np.array([4,2,2,3])
ans = [np.where(ma==i) for i in sl]

Also, if I have to resort to sorting and binary search, I will do it as a last resort (puns not intended at all sorts of levels). I am interested in finding if I'm missing something basic from the numpy library. These lists are very large, so performance is paramount.

Thanks.

Edit: Before posting I'd tried the following with dismal results:

[np.searchsorted(ma,x) for x in sl]

The solution posted by @pierre is much more performant and exactly what I was looking for.

Upvotes: 17

Views: 20588

Answers (1)

Pierre GM
Pierre GM

Reputation: 20329

Would np.searchsorted work for you ?

>>> master = np.array([1,2,3,4,5])
>>> search = np.array([4,2,2,3])
>>> np.searchsorted(master, search)
array([3, 1, 1, 2])

Upvotes: 34

Related Questions