Numpy Array: эффективно найти подходящие индексы

У меня есть два списка, один из которых огромный (миллионы элементов), другой несколько тысяч. Я хочу сделать следующее

<code>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...
</code>

Выше работает, но медленно. Есть ли способ сделать это более эффективно, не прибегая к написанию чего-либо на C?

 Michael Wild25 апр. 2012 г., 19:43
Я действительно сомневаюсь, что вы сильно ускорились бы при портировании на C, поскольку, скорее всего, операция сравнения иwhere Операция уже реализована на C.

Ответы на вопрос(3)

Вы можете использовать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]

Это должно дать вам некоторую скорость.

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]
Решение Вопроса

большого массива. Вот пример, демонстрирующий, как вы можете сократить время с ~ 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

 user135685526 апр. 2012 г., 11:25
Это именно то, что я хотел. Благодарю.
 13 янв. 2013 г., 11:51
Я не совсем уверен, но, вероятно, есть место для оптимизации с помощьюnp.searchsorted вместо петли с разделением пополам.

Ваш ответ на вопрос