在 numpy 中是否有搜索排序的字典版本？

Question

我有两个按 lex 排序的数组。

In [2]: a = np.array([1,1,1,2,2,3,5,6,6])
In [3]: b = np.array([10,20,30,5,10,100,10,30,40])
In [4]: ind = np.lexsort((b, a)) # sorts elements first by a and then by b
In [5]: print a[ind]
[1 1 1 2 2 3 5 6 6]
In [7]: print b[ind]
[ 10  20  30   5  10 100  10  30  40]

我想对 (2, 7) 和 (5, 150) 进行二分搜索，期望 (4, 7) 作为答案。

In [6]: np.lexsearchsorted((a,b), ([2, 5], [7,150]))

我们有 searchsorted 函数，但它只适用于一维数组。

Answer 1

编辑：已编辑以反映评论。

def comp_leq(t1,t2):
    if (t1[0] > t2[0]) or ((t1[0] == t2[0]) and (t1[1] > t2[1])):
        return 0
    else:
        return 1

def bin_search(L,item):
    from math import floor
    x = L[:]
    while len(x) > 1:
        index = int(floor(len(x)/2) - 1)
        #Check item
        if comp_leq(x[index], item):
            x = x[index+1:]
        else:
            x = x[:index+1]
    out = L.index(x[0])
    #If greater than all
    if item >= L[-1]:
        return len(L)
    else:
        return out

def lexsearch(a,b,items):
    z = zip(a,b)
    return [bin_search(z,item) for item in items]

if __name__ == '__main__':
    a = [1,1,1,2,2,3,5,6,6]
    b = [10,20,30,5,10,100,10,30,40]
    print lexsearch(a,b,([2,7],[5,150])) #prints [4,7]

Answer 2

这段代码似乎是为一组（完全）2 个 lexsorted 数组做的

如果您创建一组values[-1]，然后为它们创建一个带有边界的字典，您可能会更快。

我没有检查过除了发布的案例之外的其他案例，所以请确认它没有被窃听。

def lexsearchsorted_2(arrays, values, side='left'):
    assert len(arrays) == 2
    assert (np.lexsort(arrays) == range(len(arrays[0]))).all()
    # here it will be faster to work on all equal values in 'values[-1]' in one time
    boundries_l = np.searchsorted(arrays[-1], values[-1], side='left')
    boundries_r = np.searchsorted(arrays[-1], values[-1], side='right')
    # a recursive definition here will make it work for more than 2 lexsorted arrays
    return tuple([boundries_l[i] + 
                  np.searchsorted(arrays[-2[boundries_l[i]:boundries_r[i]], 
                                  values[-2][i], 
                                  side=side) 
                  for i in range(len(boundries_l))])

用法：

import numpy as np
a = np.array([1,1,1,2,2,3,5,6,6])
b = np.array([10,20,30,5,10,100,10,30,40])
lexsearchsorted_2((b, a), ([7,150], [2, 5]))  # return (4, 7)

Answer 3

我遇到了同样的问题并想出了一个不同的解决方案。您可以使用结构化数据类型将多列数据视为单个条目。结构化数据类型将允许对数据使用 argsort/sort（而不是 lexsort，尽管 lexsort 在此阶段看起来更快），然后使用标准 searchsorted。这是一个例子：

import numpy as np
from itertools import repeat

# Setup our input data
# Every row is an entry, every column what we want to sort by
# Unlike lexsort, this takes columns in decreasing priority, not increasing
a = np.array([1,1,1,2,2,3,5,6,6])
b = np.array([10,20,30,5,10,100,10,30,40])
data = np.transpose([a,b])

# Sort the data
data = data[np.lexsort(data.T[::-1])]

# Convert to a structured data-type
dt = np.dtype(zip(repeat(''), repeat(data.dtype, data.shape[1]))) # the structured dtype
data = np.ascontiguousarray(data).view(dt).squeeze(-1) # the dtype change leaves a trailing 1 dimension, ascontinguousarray is required for the dtype change

# You can also first convert to the structured data-type with the two lines above then use data.sort()/data.argsort()/np.sort(data)

# Search the data
values = np.array([(2,7),(5,150)], dtype=dt) # note: when using structured data types the rows must be a tuple
pos = np.searchsorted(data, values)

# pos is (4,7) in this example, exactly what you would want

这适用于任意数量的列，使用内置的 numpy 函数，列保持“逻辑”顺序（降低优先级），并且应该非常快。

A 在时间上比较了这两种基于 n​​umpy 的方法。

#1 是来自@j0ker5 的递归方法（下面的例子扩展了他的递归建议，并适用于任意数量的 lexsorted 行）

#2 是我的结构化数组

它们都采用相同的输入，基本上类似于searchsorted，除了a 和v 与lexsort 相同。

import numpy as np

def lexsearch1(a, v, side='left', sorter=None):
    def _recurse(a, v):
        if a.shape[1] == 0: return 0
        if a.shape[0] == 1: return a.squeeze(0).searchsorted(v.squeeze(0), side)
        bl = np.searchsorted(a[-1,:], v[-1], side='left')
        br = np.searchsorted(a[-1,:], v[-1], side='right')
        return bl + _recurse(a[:-1,bl:br], v[:-1])
    a,v = np.asarray(a), np.asarray(v)
    if v.ndim == 1: v = v[:,np.newaxis]
    assert a.ndim == 2 and v.ndim == 2 and a.shape[0] == v.shape[0] and a.shape[0] > 1
    if sorter is not None: a = a[:,sorter]
    bl = np.searchsorted(a[-1,:], v[-1,:], side='left')
    br = np.searchsorted(a[-1,:], v[-1,:], side='right')
    for i in xrange(len(bl)): bl[i] += _recurse(a[:-1,bl[i]:br[i]], v[:-1,i])
    return bl

def lexsearch2(a, v, side='left', sorter=None):
    from itertools import repeat
    a,v = np.asarray(a), np.asarray(v)
    if v.ndim == 1: v = v[:,np.newaxis]
    assert a.ndim == 2 and v.ndim == 2 and a.shape[0] == v.shape[0] and a.shape[0] > 1
    a_dt = np.dtype(zip(repeat(''), repeat(a.dtype, a.shape[0])))
    v_dt = np.dtype(zip(a_dt.names, repeat(v.dtype, a.shape[0])))
    a = np.asfortranarray(a[::-1,:]).view(a_dt).squeeze(0)
    v = np.asfortranarray(v[::-1,:]).view(v_dt).squeeze(0)
    return a.searchsorted(v, side, sorter).ravel()


a = np.random.randint(100, size=(2,10000)) # Values to sort, rows in increasing priority
v = np.random.randint(100, size=(2,10000)) # Values to search for, rows in increasing priority

sorted_idx = np.lexsort(a)
a_sorted = a[:,sorted_idx]

以及计时结果（在 iPython 中）：

# 2 rows
%timeit lexsearch1(a_sorted, v)
10 loops, best of 3: 33.4 ms per loop

%timeit lexsearch2(a_sorted, v)
100 loops, best of 3: 14 ms per loop

# 10 rows
%timeit lexsearch1(a_sorted, v)
10 loops, best of 3: 103 ms per loop

%timeit lexsearch2(a_sorted, v)
100 loops, best of 3: 14.7 ms per loop

总体而言，结构化数组方法更快，如果您将其设计为与a 和v 的翻转和转置版本一起使用，则可以更快。随着行数/键数的增加，它变得更快，从 2 行增加到 10 行时几乎没有减慢。

我没有注意到使用 a_sorted 或 a 和 sorter=sorted_idx 之间有任何显着的时间差异，所以为了清楚起见，我把它们省略了。

我相信可以使用 Cython 制作一个非常快速的方法，但这与使用纯 Python 和 numpy 一样快。