向量化搜索排序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