Numpy Array: Encuentra eficientemente los índices coincidentes

Tengo dos listas, una de las cuales es masiva (millones de elementos), la otra varios miles. Quiero hacer lo siguiente

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

Lo anterior funciona, pero es lento. ¿Hay alguna forma de hacerlo de manera más eficiente sin tener que recurrir a escribir algo en C?

En su caso, puede beneficiarse de la pre-selección de su gran variedad. Aquí está el ejemplo que muestra cómo puede reducir el tiempo de ~ 45 segundos a 2 segundos (en mi computadora portátil) (para un conjunto particular de longitudes de los arreglos 5e6 vs 1e3). Obviamente, la solución no será óptima si los tamaños de matriz serán muy diferentes. Por ejemplo, con la solución predeterminada, la complejidad es O (bigN * smallN), pero para mi solución sugerida es 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 

Salida:

Bruto 42.5278530121

No Bruto 1.57193303108

Numpy proporciona la función numpy.searchsorted: http://docs.scipy.org/doc/numpy-1.10.0/reference/generated/numpy.searchsorted.html

Ejemplo:

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

Hasta ahora no veo ninguna necesidad de adormecer; puede utilizar el defaultdict , siempre que su memoria sea suficiente, debería serlo si el número de observaciones no es demasiado grande.

 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) 

Resultado:

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

Esto debería darle algo de velocidad.