python–矢量化searchsortednumpy

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

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

解决方法:

与前一行相比,我们可以为每一行添加一些偏移量.我们将对两个数组使用相同的偏移量.这个想法是在输出数组的扁平版本之后使用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