python 矢量化搜索排序numpy

00jrzges  于 2023-04-04  发布在  Python
关注(0)|答案(3)|浏览(93)

假设我有两个数组AB,其中AB都是m x n。我现在的目标是,对于AB的每一行,找到应该在Ai行的元素插入B的相应行的位置。也就是说,我希望将np.digitizenp.searchsorted应用于AB的每一行。
我天真的解决方案是简单地遍历行。然而,这对我的应用程序来说太慢了。因此,我的问题是:有没有我还没找到的这两种算法的矢量化实现?

plicqrtu

plicqrtu1#

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

bogh5gae2#

@Divakar提供的解决方案非常适合整数数据,但要注意浮点值的精度问题,特别是当它们跨越多个数量级时(例如[[1.0, 2,0, 3.0, 1.0e+20],...])。在某些情况下,r可能非常大,以至于应用a+rb+r会擦除您尝试运行searchsorted的原始值,你只是在比较rr
为了使该方法对浮点数据更健壮,您可以将行信息作为值的一部分(作为结构化数据类型)嵌入到数组中,并在这些结构化数据类型上运行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
ie3xauqp

ie3xauqp3#

我认为@Divakar的解决方案需要以下两个注解(由于声誉原因,我无法添加评论):
1.这是一个计算(向量化)优化,不是算法优化。这意味着在理论意义上没有复杂度增益。
1.事实上,在某些情况下,它会,有时会慢10倍以上。例如,当使用形状(20, 10000)而不是(10000,20)时,当b.shape[1]a.shape[1]小得多时,或者当有太多数据时,在一条很长的松散线上工作时内存效率低下。

相关问题