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

Dav*_*ave 12 python arrays numpy

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

我想快速返回a1中所有元素的a2索引.这种时间密集的方式显然是:

from operator import indexOf
indices = []
for i in a1:
    indices.append(indexOf(a2,i))
Run Code Online (Sandbox Code Playgroud)

这当然需要很长时间,其中a2很大.我也可以使用numpy.where()(尽管a1中的每个条目在a2中只出现一次),但我不相信它会更快.我也可以只遍历一次大阵列:

for i in xrange(len(a2)):
    if a2[i] in a1:
        indices.append(i)
Run Code Online (Sandbox Code Playgroud)

但我确信有一种更快,更"笨拙"的方式 - 我查看了numpy方法列表,但找不到合适的东西.

提前谢谢了,

d

Alo*_*hal 12

怎么样

numpy.nonzero(numpy.in1d(a2, a1))[0]
Run Code Online (Sandbox Code Playgroud)

这应该很快.从我的基本测试,它比你的第二个代码片段快了约7倍len(a2) == 100,len(a1) == 10000且只有一个索引45共同元素此假定都a1a2有没有重复的元素.

  • 对于阅读此内容的任何人:似乎`setmember1d`自numpy 1.4以来已被重命名为`in1d`. (3认同)