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

Question

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

例如

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

我想找到a中每个元素在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 给出的 here ，使用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]