向量化搜索排序的 numpy

Question

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

我天真的解决方案是简单地遍历行。但是，这对于我的应用程序来说太慢了。因此，我的问题是：是否存在我无法找到的任一算法的矢量化实现？

Answer 1

@Divakar 提供的解决方案非常适合整数数据，但要注意浮点值的精度问题，尤其是当它们跨越多个数量级时（例如 [[1.0, 2,0, 3.0, 1.0e+20],...]）。在某些情况下，r 可能太大，以至于应用a+r 和b+r 会消除您尝试运行searchsorted 的原始值，而您只是将r 与r 进行比较。

为了使该方法对浮点数据更加稳健，您可以将行信息作为值的一部分（作为结构化 dtype）嵌入到数组中，然后对这些结构化 dtype 运行 searchsorted。

def searchsorted_2d (a, v, side='left', sorter=None):
  import numpy as np

  # Make sure a and v are numpy arrays.
  a = np.asarray(a)
  v = np.asarray(v)

  # Augment a with row id
  ai = np.empty(a.shape,dtype=[('row',int),('value',a.dtype)])
  ai['row'] = np.arange(a.shape[0]).reshape(-1,1)
  ai['value'] = a

  # Augment v with row id
  vi = np.empty(v.shape,dtype=[('row',int),('value',v.dtype)])
  vi['row'] = np.arange(v.shape[0]).reshape(-1,1)
  vi['value'] = v

  # Perform searchsorted on augmented array.
  # The row information is embedded in the values, so only the equivalent rows 
  # between a and v are considered.
  result = np.searchsorted(ai.flatten(),vi.flatten(), side=side, sorter=sorter)

  # Restore the original shape, decode the searchsorted indices so they apply to the original data.
  result = result.reshape(vi.shape) - vi['row']*a.shape[1]

  return result

编辑：这种方法的时机太糟糕了！

In [21]: %timeit searchsorted_2d(a,b)
10 loops, best of 3: 92.5 ms per loop

你最好只在数组上使用map：

In [22]: %timeit np.array(list(map(np.searchsorted,a,b)))
100 loops, best of 3: 13.8 ms per loop

对于整数数据，@Divakar 的方法仍然是最快的：

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

Answer 2

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