矢量化搜索排序numpy

Question

假设我有两个数组A和B,其中两个A和B的m x n.我的目标是现在,对于每一行A和B,找到我应该在哪里插入行的元素i中A的相应行中B.也就是说,我希望申请np.digitize或np.searchsorted每行A和B.

我天真的解决方案是简单地迭代行.但是,这对我的应用来说太慢了.因此,我的问题是:是否存在我无法找到的算法的矢量化实现？

Answer 1

与前一行相比,我们可以为每一行添加一些偏移量.我们将对两个数组使用相同的偏移量.这个想法是np.searchsorted在此后使用平面版本的输入数组,因此每一行都b将被限制在相应的行中查找排序位置a.此外,为了使其适用于负数,我们也需要抵消最小数量.

所以,我们将有一个像这样的矢量化实现 -

def searchsorted2d(a,b):
    m,n = a.shape
    max_num = np.maximum(a.max() - a.min(), b.max() - b.min()) + 1
    r = max_num*np.arange(a.shape[0])[:,None]
    p = np.searchsorted( (a+r).ravel(), (b+r).ravel() ).reshape(m,-1)
    return p - n*(np.arange(m)[:,None])

运行时测试 -

In [173]: def searchsorted2d_loopy(a,b):
     ...:     out = np.zeros(a.shape,dtype=int)
     ...:     for i in range(len(a)):
     ...:         out[i] = np.searchsorted(a[i],b[i])
     ...:     return out
     ...: 

In [174]: # Setup input arrays
     ...: a = np.random.randint(11,99,(10000,20))
     ...: b = np.random.randint(11,99,(10000,20))
     ...: a = np.sort(a,1)
     ...: b = np.sort(b,1)
     ...: 

In [175]: np.allclose(searchsorted2d(a,b),searchsorted2d_loopy(a,b))
Out[175]: True

In [176]: %timeit searchsorted2d_loopy(a,b)
10 loops, best of 3: 28.6 ms per loop

In [177]: %timeit searchsorted2d(a,b)
100 loops, best of 3: 13.7 ms per loop