一次获取 NumPy 数组中多个元素的索引

Question

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

Eg

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

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.

Answer 5

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)]