一次获取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]