在numpy中是否有类似于searchsorted的字典排序版本？

Question

在numpy中是否有类似于searchsorted的字典排序版本？

3

我有两个按字典序排序的数组。

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函数，但它仅适用于1D数组。

- pvncad

你能提供更多关于你尝试做什么的信息吗？这个在numpy中的实现对我来说很奇怪。 - tsvikas

我有两组时间序列，它们都由相同的键进行索引。类似于（用户、时间、位置）和（用户、时间、购买）在两个不同的地方。我正在尝试将它们合并起来，以获得类似于（用户、时间、位置、购买）的东西。 - pvncad

3个回答

1

这段代码似乎可以对一组（恰好）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)

- tsvikas

感谢您的解决方案。其中一个问题是它会创建两个新数组。 - pvncad

是的，但它们的大小与“values”的大小相同，因此对于大的“a”和“b”仍然有效。 - tsvikas

1

我遇到了相同的问题，并提出了另一种解决方案。您可以使用结构化数据类型将多列数据视为单个条目。结构化数据类型将允许在数据上使用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函数，列保持“逻辑”顺序（优先级降低），且速度应该很快。

我比较了这两种基于numpy的方法的时间。

#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的最快方法了。

- coderforlife

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Ross B. · Accepted Answer

编辑： 根据评论进行了编辑。

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]