Question

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

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

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

Answer 2

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

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

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

矢量化搜索排序numpy

2 个答案: