返回两个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

在编程中，有时候我们会遇到一些问题，特别是在使用某些工具或库的时候。这些问题可能会让我们感到困惑，不知道该怎么解决。比如，有人可能在使用某个特定的功能时，发现它没有按照预期的方式工作。这种情况下，通常我们会去查找相关的资料，看看有没有人遇到过类似的问题，或者有没有解决方案。

在StackOverflow上，很多人会分享他们的经验和解决方案。你可以看到各种各样的问题和答案，有些是关于代码的，有些是关于使用工具的。通过这些讨论，我们可以学习到很多实用的知识，帮助我们更好地理解编程的世界。

总之，遇到问题时，不要害怕去寻求帮助，社区里有很多人愿意分享他们的经验和解决方法，这对我们学习编程非常有帮助。

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

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

顺便说一下，我对 numpy 不太了解，所以可能还有更“符合 numpy 风格”的做法，但这是我在纯 Python 中的做法。

Answer 3

怎么样呢

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

这个应该很快。根据我简单的测试，对于 len(a2) == 100 和 len(a1) == 10000 的情况，并且只有一个共同的元素在索引45的位置，这个代码比你第二段代码快大约7倍。这个测试是基于 a1 和 a2 中没有重复元素的假设。