返回两个 numpy 数组之间的公共元素索引答案

【问题标题】：Return common element indices between two numpy arrays返回两个 numpy 数组之间的公共元素索引
【发布时间】：2011-01-20 23:22:12
【问题描述】：

我有两个数组，a1 和 a2。假设len(a2) >> len(a1)，a1 是 a2 的子集。

我想要一种快速的方法来返回 a1 中所有元素的 a2 索引。执行此操作的时间密集型方法显然是：

from operator import indexOf
indices = []
for i in a1:
    indices.append(indexOf(a2,i))

这当然需要很长时间 a2 很大。我也可以使用 numpy.where() 代替（尽管 a1 中的每个条目在 a2 中只会出现一次），但我不相信它会更快。我也可以只遍历一次大数组：

for i in xrange(len(a2)):
    if a2[i] in a1:
        indices.append(i)

但我确信有一种更快、更“numpy”的方法 - 我查看了 numpy 方法列表，但找不到任何合适的方法。

非常感谢，

D

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

这些方法对我来说都很慢。以下方法做得很快。索引列表具有第一个列表中第二个列表中常见的元素的索引。

index=[]
d={}
for j in range(len(first_list)):
    name=first_list[j]
    d[name]=j
    
for i in range(len(second_list)):
    name=second_list[i]
    index.append(d[name])

【讨论】：

【解决方案2】：

与@AlokSinghal 非常相似，但你得到的是一个已经扁平化的版本。

numpy.flatnonzero(numpy.in1d(a2, a1))

【讨论】：

【解决方案3】：

numpy_indexed 包（免责声明：我是它的作者）包含 list.index 的矢量化等效项；性能应该类似于当前接受的答案，但作为奖励，它还可以让您使用“缺失”kwarg 明确控制缺失值。

import numpy_indexed as npi
indices = npi.indices(a2, a1, missing='raise')

此外，它还适用于多维数组，即在另一组中查找一组行的索引。

【讨论】：

【解决方案4】：

怎么样

numpy.nonzero(numpy.in1d(a2, a1))[0]

这应该很快。根据我的基本测试，对于len(a2) == 100、len(a1) == 10000 和索引 45 处只有一个公共元素，它比您的第二个代码 sn-p 快 7 倍左右。这假设 a1 和 a2 都没有重复元素.

【讨论】：

我将您的解决方案与上述 Dave Kirby 的解决方案进行了比较，对于 len(a2) == 12347424、len(a1) == 1338，这个解决方案大约快 1.35 倍，所以这个解决方案得到了我的投票 - 谢谢!
对于任何阅读本文的人：似乎 setmember1d 自 numpy 1.4 起已重命名为 in1d。
@AlokSinghal 感谢您的关注！
@DanielWatkins 由于 numpy 1.4 现在已经很老了，我已经更新了我的答案以使用 in1d。

【解决方案5】：

index = in1d(a2,a1)
result = a2[index]

【讨论】：

【解决方案6】：

怎么样：

wanted = set(a1)
indices =[idx for (idx, value) in enumerate(a2) if value in wanted]

这应该是 O(len(a1)+len(a2)) 而不是 O(len(a1)*len(a2))

注意，我不知道 numpy，所以可能有一种更“numpythonic”的方式来做，但这是我在纯 python 中的做法。

【讨论】：

应该是 enumerate(a2) 吗？