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

Question

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

例如

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

我想在b中找到a的每个元素的索引，即： [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.flatnonzero与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

在传递这个相关问题的答案时，有很多方法可以一次获取多个项目的索引： Is there a NumPy function to return the first index of something in an array？ . 答案的多样性和创造性表明没有单一的最佳实践，因此如果您的上述代码有效且易于理解，我会说保留它。

我个人发现这种方法既高效又易于阅读： https ://stackoverflow.com/a/23994923/3823857

为您的示例进行调整：

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]

我个人喜欢添加一点错误处理，以防b a不存在。

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]

对于我的用例，它非常快，因为它依赖于快速的 Python 部分（列表推导、.index()、集合、numpy 索引）。 仍然希望看到与 VLOOKUP 等效的 NumPy 甚至是 Pandas 合并的东西。 但这似乎暂时有效。

Answer 5

这里的所有解决方案都建议使用线性搜索。 您可以使用np.argsort和np.searchsorted来显着加快大型数组的速度：

sorter = b.argsort()
i = sorter[np.searchsorted(b, a, sorter=sorter)]