使用NumPy进行大型数组搜索

Question

我有两个整数数组

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

其中 len(a) 约为15,000，len(b) 约为2百万。

我想要找到数组 b 中与数组 a 相匹配的元素的索引。现在，我使用列表推导式和 numpy.argwhere() 来实现这个目标:

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

然而，显然，完成这个任务需要很长时间。而且数组a会变得更大，所以这不是一个明智的选择。

考虑到我正在处理的大型数组，是否有更好的方法来实现这个结果？目前这需要大约5分钟的时间。需要任何加速！

更多信息：我希望索引也能匹配数组a的顺序。(感谢Charles)

- Carl M

---

也许你可以创建一个哈希映射，将元素从 a 映射到它们各自的索引。然后你只需要在映射中查找它们。 - tobias_k

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

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

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

- tobias_k

---

我相信这正是我需要的！我想这就是你所说的哈希表吧？非常感谢你的时间。 - Carl M

是的，哈希表或多或少可以理解为字典的另一个词。在Python中，它被称为字典或dict，或者只是{}；在Java中，它是一个Map或HashMap。对于混淆造成的困扰，我们表示歉意。 - tobias_k

你有什么想法可以在b中找不到a的项时添加故障保护？ - Carl M

我刚刚创建了一个独立的集合，并使用了以下代码：[ indices[c] if c in b_set else -99 for c in a ] - Carl M

你不需要一个单独的 set。在 dict 中查找也是 O(1) 的，所以你可以直接使用 indices[c] if c in indices else -99，或者使用带有默认值的 get 方法，即 indices.get(e, -99)。 - tobias_k

哦，太棒了。谢谢 @tobias_k！ - Carl M

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

- Charles

---

谢谢！不过，我希望索引的顺序与a中的元素顺序相同。这样可以使它们与b中的顺序一致。明白我的意思吗？ - Carl M

是的，但您应该澄清您的问题，因为那不够清楚。 - Charles