поиск отсортированных элементов в отсортированной последовательности

Я хочу найти последовательность элементов в отсортированном массиве значений. Я знаю, что с numpy я могу сделать:

l = np.searchsorted(values, items)

Это имеет сложность O (len (items) * log (len (values))). Однако мои объекты также сортируются, поэтому я могу сделать это в O (len (items) + len (values)):

 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

Проблема в том, что эта версия в чистом питоне медленнее, чем searchsorted из-за цикла python, даже для больших len (элементов) и len (значений) (~ 10 ^ 6).

Любая идея, как «векторизовать» этот цикл с numpy?

One Solution collect form web for “поиск отсортированных элементов в отсортированной последовательности”

Некоторые примеры данных:

 # 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()

Исходный фрагмент кода, а также реализация предложения @ 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

Тест на правильность к наивному np.searchsorted :

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

Тайминги:

 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

Таким образом, для входов такого размера наивное использование np.searchsorted удобно превосходит ваш исходный код, а также предложение PeterE.

Обновить

Чтобы избежать каких-либо эффектов кеширования, которые могут исказить тайминги, мы можем создать новый набор случайных входных массивов для каждой итерации эталона:

 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

Обновление 2

Не так сложно написать функцию Cython, которая будет бить np.searchsorted для случая, когда будут отсортированы как values и 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

Проверка правильности:

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

Ориентир:

 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

Однако улучшение производительности довольно незначительно. Если это не является серьезным узким местом в вашем коде, вы должны, вероятно, придерживаться np.searchsorted .