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]

Answer 2

This is a simple one-liner using the numpy-indexed package (disclaimer: I am its author):

import numpy_indexed as npi
idx = npi.indices(b, a)

The implementation is fully vectorized, and it gives you control over the handling of missing values. Moreover, it works for nd-arrays as well (for instance, finding the indices of rows of a in b).

Answer 3

For an order-agnostic solution, you can use np.flatnonzero with np.isin (v 1.13+).

import numpy as np

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

res = np.flatnonzero(np.isin(a, b))  # NumPy v1.13+
res = np.flatnonzero(np.in1d(a, b))  # earlier versions

# array([0, 1, 2], dtype=int64)

Answer 4

All of the solutions here recommend using a linear search. You can use np.argsort and np.searchsorted to speed things up dramatically for large arrays:

sorter = b.argsort()
i = sorter[np.searchsorted(b, a, sorter=sorter)]

Answer 5

There are a bunch of approaches for getting the index of multiple items at once mentioned in passing in answers to this related question: Is there a NumPy function to return the first index of something in an array?. The wide variety and creativity of the answers suggests there is no single best practice, so if your code above works and is easy to understand, I'd say keep it.

I personally found this approach to be both performant and easy to read: https://stackoverflow.com/a/23994923/3823857

Adapting it for your example:

import numpy as np

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

indices = [b_list.index(x) for x in a]
vals_at_indices = b_array[indices]

I personally like adding a little bit of error handling in case a value in a does not exist in b.

import numpy as np

a = np.array([1, 2, 4])
b_list = [1, 2, 3, 10, 4]
b_array = np.array(b_list)
b_set = set(b_list)

indices = [b_list.index(x) if x in b_set else np.nan for x in a]
vals_at_indices = b_array[indices]

For my use case, it's pretty fast, since it relies on parts of Python that are fast (list comprehensions, .index(), sets, numpy indexing). Would still love to see something that's a NumPy equivalent to VLOOKUP, or even a Pandas merge. But this seems to work for now.