将已排序的项目搜索到已排序的序列中

问题描述

我想在值的排序数组中找到一系列项目.我知道使用numpy我可以做到:

I want to find a sequence of items in a sorted array of values.I know that with numpy I can do:

l = np.searchsorted(values, items)

这具有O(len(items)* log(len(values)))的复杂度.但是，我的项目也已排序，因此我可以在O(len(items)+ len(values))中进行以下操作:

This has the complexity of O(len(items)*log(len(values))).However, my items are also sorted, so I can do it in O(len(items)+len(values)) doing:

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

问题在于，纯python的此版本由于python循环而比searchsorted慢得多，即使对于大len(items)和len(values)(〜10 ^ 6)也是如此.

The problem is that this version in pure python is way slower than searchsorted because of the python loop, even for large len(items) and len(values) (~10^6).

有什么想法如何使用numpy对向量循环进行矢量化"处理吗?

Any idea how to "vectorize" this loop with numpy?

推荐答案

一些示例数据:

# 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的建议的实现:

Your original code snippet as well as an implementation of @PeterE's suggestion:

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的正确性:

Test for correctness against naive 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的建议.

So for inputs of this size, naive use of np.searchsorted handily beats your original code, as well as PeterE's suggestion.

为避免任何可能会影响时序的缓存影响，我们可以为基准测试的每次迭代生成一组新的随机输入数组:

To avoid any caching effects that might skew the timings, we can generate a new set of random input arrays for each iteration of the benchmark:

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

编写同时对values和items进行排序的Cython函数打败np.searchsorted并不难.

Update 2

It's not that hard to write a Cython function that will beat np.searchsorted for the case where both values and items are sorted.

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.

The performance improvement is fairly marginal, though. Unless this is a serious bottleneck in your code then you should probably stick with np.searchsorted.

这篇关于将已排序的项目搜索到已排序的序列中的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！