buscar elementos ordenados en una secuencia ordenada

Quiero encontrar una secuencia de elementos en una matriz ordenada de valores. Sé que con numpy puedo hacer:

l = np.searchsorted(values, items) 

Esto tiene la complejidad de O (len (elementos) * log (len (valores))). Sin embargo, mis artículos también están ordenados, así que puedo hacerlo en O (len (items) + len (valores)) haciendo:

 l = np.zeros(items.size, dtype=np.int32) k, K = 0, len(values) for i in range(len(items)): while k < K and values[k] < items[i]: k += 1 l[i] = k 

El problema es que esta versión en python puro es mucho más lenta que la búsqueda debido al bucle de python, incluso para grandes len (elementos) y len (valores) (~ 10 ^ 6).

¿Alguna idea de cómo “vectorizar” este bucle con numpy?

Algunos datos de ejemplo:

 # some example data np.random.seed(0) n_values = 1000000 n_items = 100000 values = np.random.rand(n_values) items = np.random.rand(n_items) values.sort() items.sort() 

Su fragmento de código original, así como una implementación de la sugerencia de @ PeterE:

 def original(values, items): l = np.empty(items.size, dtype=np.int32) k, K = 0, len(values) for i, item in enumerate(items): while k < K and values[k] < item: k += 1 l[i] = k return l def peter_e(values, items): l = np.empty(items.size, dtype=np.int32) last_idx = 0 for i, item in enumerate(items): last_idx += values[last_idx:].searchsorted(item) l[i] = last_idx return l 

Prueba de corrección contra naive np.searchsorted :

 ss = values.searchsorted(items) print(all(original(values, items) == ss)) # True print(all(peter_e(values, items) == ss)) # True 

Tiempos:

 In [1]: %timeit original(values, items) 10 loops, best of 3: 115 ms per loop In [2]: %timeit peter_e(values, items) 10 loops, best of 3: 79.8 ms per loop In [3]: %timeit values.searchsorted(items) 100 loops, best of 3: 4.09 ms per loop 

Por lo tanto, para entradas de este tamaño, el uso ingenuo de np.searchsorted supera fácilmente su código original, así como la sugerencia de PeterE.

Actualizar

Para evitar cualquier efecto de almacenamiento en caché que pueda sesgar los tiempos, podemos generar un nuevo conjunto de matrices de entrada aleatorias para cada iteración del punto de referencia:

 In [1]: %%timeit values = np.random.randn(n_values); items = np.random.randn(n_items); values.sort(); items.sort(); original(values, items) .....: 10 loops, best of 3: 115 ms per loop In [2]: %%timeit values = np.random.randn(n_values); items = np.random.randn(n_items); values.sort(); items.sort(); peter_e(values, items) .....: 10 loops, best of 3: 79.9 ms per loop In [3]: %%timeit values = np.random.randn(n_values); items = np.random.randn(n_items); values.sort(); items.sort(); values.searchsorted(items) .....: 100 loops, best of 3: 4.08 ms per loop 

Actualización 2

No es tan difícil escribir una función de Cython que np.searchsorted a np.searchsorted en el caso de que se np.searchsorted tanto los values como los items .

search_doubly_sorted.pyx :

 import numpy as np cimport numpy as np cimport cython @cython.boundscheck(False) @cython.wraparound(False) def search_doubly_sorted(values, items): cdef: double[:] _values = values.astype(np.double) double[:] _items = items.astype(np.double) long n_items = items.shape[0] long n_values = values.shape[0] long[:] out = np.empty(n_items, dtype=np.int64) long ii, jj, last_idx last_idx = 0 for ii in range(n_items): for jj in range(last_idx, n_values): if _items[ii] <= _values[jj]: break last_idx = jj out[ii] = last_idx return out.base 

Prueba de corrección:

 In [1]: from search_doubly_sorted import search_doubly_sorted In [2]: print(all(search_doubly_sorted(values, items) == values.searchsorted(items))) # True 

Punto de referencia:

 In [3]: %timeit values.searchsorted(items) 100 loops, best of 3: 4.07 ms per loop In [4]: %timeit search_doubly_sorted(values, items) 1000 loops, best of 3: 1.44 ms per loop 

Sin embargo, la mejora del rendimiento es bastante marginal. A menos que esto sea un grave cuello de botella en su código, entonces probablemente debería quedarse con np.searchsorted .