Numpy数组：高效查找匹配的索引

Question

我有两个列表，其中一个非常大（数百万个元素），另一个只有几千个。我想要做以下操作：

bigArray=[0,1,0,2,3,2,,.....]

smallArray=[0,1,2,3,4]

for i in len(smallArray):
  pts=np.where(bigArray==smallArray[i])
  #Do stuff with pts...

上述方法可行，但速度较慢。有没有更高效的方法来完成这个任务，而不需要编写C代码？

- user1356855

---

我真的怀疑，如果将其移植到C语言中，你不太可能获得更快的速度，因为很可能比较操作和“where”操作已经在C语言中实现了。 - Michael Wild

Answer 1

Numpy提供了numpy.searchsorted函数：http://docs.scipy.org/doc/numpy-1.10.0/reference/generated/numpy.searchsorted.html 示例：

>>> import numpy as np
>>> sorted = np.argsort(big_list)
>>> r = np.searchsorted(big_list, small_list, side='right',sorter=sorted)
>>> l  = np.searchsorted(big_list, small_list, side='left',sorter=sorted)
>>> for b, e in zip(l, r):
...     inds = sorted[b:e]

- Raphael Campos

---

我没有测试二分法版本，但在我的快速实验中，这至少比defaultdic查找答案更快。在我的设置中，速度比默认方法快了多达2倍。 - Fabricio

Answer 2

在你的情况下，对大型数组进行预排序可能会有所帮助。以下是一个示例，演示如何将时间从约45秒减少到2秒（在我的笔记本电脑上）（针对一组特定的数组长度5e6与1e3）。显然，如果数组大小差异很大，该解决方案将不会是最优的。例如，对于默认解决方案，复杂度为O(bigN*smallN)，但对于我建议的解决方案，它为O((bigN+smallN)*log(bigN))。

import numpy as np, numpy.random as nprand, time, bisect

bigN = 5e6
smallN = 1000
maxn = 1e7
nprand.seed(1)  
bigArr = nprand.randint(0, maxn, size=bigN)
smallArr = nprand.randint(0, maxn, size=smallN)

# brute force 
t1 = time.time()
for i in range(len(smallArr)):
    inds = np.where(bigArr == smallArr[i])[0]
t2 = time.time()
print "Brute", t2-t1

# not brute force (like nested loop with index scan)
t1 = time.time()
sortedind = np.argsort(bigArr)
sortedbigArr = bigArr[sortedind]
for i in range(len(smallArr)):
    i1 = bisect.bisect_left(sortedbigArr, smallArr[i])
    i2 = bisect.bisect_right(sortedbigArr, smallArr[i])
    inds = sortedind[i1:i2]
t2=time.time()
print "Non-brute", t2-t1

输出:

暴力破解 42.5278530121

非暴力破解 1.57193303108

- sega_sai

---

我不是完全确定，但可能可以通过使用np.searchsorted替代二分循环来进行优化。 - Michael

Answer 3

目前我还没有看到需要使用numpy的必要性；你可以利用defaultdict，只要你的内存足够，如果观测数量不是太多（不超过数百万），那么就可以使用。

big_list = [0,1,0,2,3,2,5,6,7,5,6,4,5,3,4,3,5,6,5]
small_list = [0,1,2,3,4]

from collections import defaultdict

dicto = defaultdict(list) #dictionary stores all the relevant coordinates
                          #so you don't have to search for them later

for ind, ele in enumerate(big_list):
    dicto[ele].append(ind)

结果：

>>> for ele in small_list:
...     print dicto[ele]
... 
[0, 2]
[1]
[3, 5]
[4, 13, 15]
[11, 14]

这应该会让你更快。

- Akavall