向量化搜索排序的numpy

向量化搜索排序的numpy

与上一行相比，我们可以为每行添加一些偏移量。我们将对两个数组使用相同的偏移量。想法是

np.searchsorted

在此后在输入数组的扁平版本上使用，因此from中的每一行都

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]: TrueIn [176]: %timeit searchsorted2d_loopy(a,b)10 loops, best of 3: 28.6 ms per loopIn [177]: %timeit searchsorted2d(a,b)100 loops, best of 3: 13.7 ms per loop