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

Question

我有两个数组，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

Answer 1

怎么样

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

这应该很快。从我的基本测试开始，它比第len(a2) == 100，len(a1) == 10000的第二个代码段快7倍，而索引45只有一个公共元素。这假设a1和{{1没有重复元素。

Answer 2

怎么样：

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

NB我不知道numpy所以可能有一种更“'numpythonic'的方式来实现它，但这就是我在纯python中做的方式。

Answer 3

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

Answer 4

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

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

Answer 5

numpy_indexed包（免责声明：我是它的作者）包含了vector.index的向量化等价物;性能应该与当前接受的答案类似，但作为奖励，它使您可以使用“缺失”来明确控制缺失值。 kwarg。

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

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