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

Question

有没有办法一次获取NumPy数组中的几个元素的索引？

例如

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

我想找到ain 中每个元素的索引b,即:[0,1,4].

我发现我使用的解决方案有点冗长:

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

输出:

c: [0 1 4]

Answer 1

你可以使用in1d和nonzero(或者where就此而言):

>>> np.in1d(b, a).nonzero()[0]
array([0, 1, 4])

这适用于您的示例数组,但通常返回的索引数组不符合值的顺序a.这可能是一个问题,具体取决于您下一步要做什么.

在这种情况下,更好的答案是一个@Jaime给出了这里,使用searchsorted:

>>> sorter = np.argsort(b)
>>> sorter[np.searchsorted(b, a, sorter=sorter)]
array([0, 1, 4])

这将返回它们出现时的值的索引a.例如:

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

这是一个简单的单行使用numpy-indexed包(免责声明:我是它的作者):

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

该实现是完全矢量化的,它使您可以控制缺失值的处理.此外,它也适用于nd数组(例如,查找b中a的行的索引).

Answer 3

对于与顺序无关的解决方案，您可以使用np.flatnonzerowith 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)