Question

我有两个整数数组

a = numpy.array([1109830922873, 2838383, 839839393, ..., 29839933982])
b = numpy.array([2838383, 555555555, 2839474582, ..., 29839933982])

其中len(a) ~15,000和len(b) ~2百万。

我想要的是找到与数组a中的数组匹配的数组b元素的索引。现在，我使用列表理解和numpy.argwhere()来实现这一目标：

bInds = [ numpy.argwhere(b == c)[0] for c in a ]

然而，很明显，这需要很长时间才能完成。阵列a也会变大，所以这不是一条明智的选择。

考虑到我在这里处理的大型阵列，是否有更好的方法来实现这一结果？目前大约需要5分钟左右。需要加速！

更多信息：我希望索引也匹配数组a的顺序。（谢谢查尔斯）

Answer 1

除非我弄错了，否则您的方法会一次又一次地搜索整个数组b以查找a的每个元素。

或者，您可以创建一个字典，将单个元素从b映射到它们的索引。

indices = {}
for i, e in enumerate(b):
    indices[e] = i                      # if elements in b are unique
    indices.setdefault(e, []).append(i) # otherwise, use lists

然后，您可以使用此映射快速查找可在a中找到b元素的索引。

bInds = [ indices[c] for c in a ]

Answer 2

这需要大约一秒钟才能运行。

import numpy

#make some fake data...
a = (numpy.random.random(15000) * 2**16).astype(int)
b = (numpy.random.random(2000000) * 2**16).astype(int)

#find indcies of b that are contained in a.
set_a = set(a)
result = set()
for i,val in enumerate(b):
    if val in set_a:
        result.add(i)

result = numpy.array(list(result))
result.sort()

print result

用numpy搜索大型数组

2 个答案: