假设我有两个数组A和B,其中A和B都是m x n。我现在的目标是,对于A和B的每一行,找到应该在A的i行的元素插入B的相应行的位置。也就是说,我希望将np.digitize或np.searchsorted应用于A和B的每一行。 我天真的解决方案是简单地遍历行。然而,这对我的应用程序来说太慢了。因此,我的问题是:有没有我还没找到的这两种算法的矢量化实现?
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
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
3条答案
按热度按时间plicqrtu1#
我们可以为每一行添加一些与前一行相比的偏移量。我们将对两个数组使用相同的偏移量。我们的想法是在此后的输入数组的扁平化版本上使用
np.searchsorted
,因此b
中的每一行将被限制为在a
中找到相应行中的排序位置。此外,为了使其也适用于负数,我们只需要对最小数量进行偏移。所以,我们会有一个矢量化的实现,就像这样-
运行时间测试-
bogh5gae2#
@Divakar提供的解决方案非常适合整数数据,但要注意浮点值的精度问题,特别是当它们跨越多个数量级时(例如
[[1.0, 2,0, 3.0, 1.0e+20],...]
)。在某些情况下,r
可能非常大,以至于应用a+r
和b+r
会擦除您尝试运行searchsorted
的原始值,你只是在比较r
和r
。为了使该方法对浮点数据更健壮,您可以将行信息作为值的一部分(作为结构化数据类型)嵌入到数组中,并在这些结构化数据类型上运行searchsorted。
**编辑:**这种方法的时机是糟糕透顶!
你最好只使用
map
而不是数组:对于整数数据,@Divakar的方法仍然是最快的:
ie3xauqp3#
我认为@Divakar的解决方案需要以下两个注解(由于声誉原因,我无法添加评论):
1.这是一个计算(向量化)优化,不是算法优化。这意味着在理论意义上没有复杂度增益。
1.事实上,在某些情况下,它会慢,有时会慢10倍以上。例如,当使用形状
(20, 10000)
而不是(10000,20)
时,当b.shape[1]
比a.shape[1]
小得多时,或者当有太多数据时,在一条很长的松散线上工作时内存效率低下。