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 

Выход:

Brute +42,5278530121

Non-brute 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]