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

为君一笑 提交于 2019-11-26 08:13:21

问题


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]

回答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]



回答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).




回答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)



回答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.



来源:https://stackoverflow.com/questions/32191029/getting-the-indices-of-several-elements-in-a-numpy-array-at-once

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!