Numpy Array：高效地找到匹配的索引

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寫東西？

Answer 1

在你的情況下，你可能會從預測你的大陣列中受益。下面是演示如何將時間從約45秒減少到2秒（在我的筆記本電腦上）（針對陣列5e6和1e3的一組特定長度）的示例。顯然，如果陣列大小差別很大，解決方案將不會最佳。例如。使用默認的解決方案的複雜度爲O（bigN * smallN），但我建議的解決方案是O（（bigN + smallN）*日誌（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

Answer 2

到目前爲止，我沒有看到任何需要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] 

這應該給你一些速度。

Answer 3

與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]