高效的一维数组比较、缩放和求和

Question

我正面临性能问题，因为一段 python 代码应该执行以下操作：

我有 2 个数组 A 和 B 的值未排序，我想构造一个新数组 C，它包含以下每个索引 i：

C[i]= sum(flag*B[k] for k so that flag = A[k]<=A[i])

我有两种方法：

1) 非常直接的方式：

M = len(A)
C = np.zeros(M)
for i in xrange(M):
    value = A[i]
    flag = A <= value
    C[i] = np.sum(flag * B)

2) 尝试使用 numpy 排序功能：

indices_sorted = np.argsort(A)
C_sort = np.zeros(M)
for i in xrange(M):
    index = np.where(indices_sorted==i)
    for k in xrange(index[0][0]+1):
        C_sort[i] += B[indices_sorted[k]]

结果是，对于 5000 个元素的数组，第一个更快（因子 40-50）。

我没想到第二次会那么糟糕，第一次尝试也不够快......

你们能给我一个更好的方法吗？

提前致谢。

Answer 1

假设A 和B 是相同形状的一维数组，您可以通过将A 扩展为二维数组然后进行比较来使用broadcasting，从而基本上将每个元素与每个元素进行比较其他元素以矢量化方式。然后，使用B 执行元素乘法，broadcasting 再次发挥作用。最后沿第二个轴求和以获得最终输出。实现看起来像这样 -

C = ((A <= A[:,None])*B).sum(1)

您可以使用np.dot 模拟elementwise multiplication and summing 和matrix-multiplication 的相同行为，以获得更有效的解决方案，就像这样 -

C = (A <= A[:,None]).dot(B)

这是另一种基于np.take 索引和np.bincount 计数的方法-

row,col = np.nonzero(A <= A[:,None])
C = np.bincount(row,np.take(B,col))

对于巨大的数据大小，创建 2D 掩码 (A <= A[:,None] 的内存开销可能会抵消性能。因此，作为对现有循环代码的优化，您可以引入matrix-multiplication 来替换元素乘法和求和。因此，np.sum(flag * B) 可以替换为flag.dot(B)。引入一些其他优化技巧，您将拥有像这样的修改版本 -

M = len(A)
C = np.empty(M)
for i in xrange(M):
    C[i] = (A <= A[i]).dot(B)

终于！这是np.cumsum的获胜者-

idx = A.argsort()
C = B[idx].cumsum()[idx.argsort()]

下面是关于它如何工作以及为什么工作的快速解释：

您正在执行逐元素比较，然后根据比较结果对 B 中的元素求和。现在，如果A 是一个排序数组，那么输出C 基本上就是cumsum 版本的B。因此，对于一般未排序的情况，您需要对B 进行A 的argsort 排序，对其执行cumsum，最后根据原始未排序的顺序重新排列元素。

---

运行时测试

定义方法 -

def org_app(A,B):
    M = len(A)
    C = np.zeros(M)
    for i in range(M):
        value = A[i]
        flag = A <= value
        C[i] = np.sum(flag * B)
    return C

def sum_based(A,B):
    return ((A <= A[:,None])*B).sum(1)

def dot_based(A,B):
    return (A <= A[:,None]).dot(B)

def bincount_based(A,B):
    row,col = np.nonzero(A <= A[:,None])
    return np.bincount(row,np.take(B,col))

def org_app_modified(A,B):
    M = len(A)
    C = np.empty(M)
    for i in xrange(M):
        C[i] = (A <= A[i]).dot(B)
    return C

def cumsum_trick(A,B):
    idx = A.argsort()
    return B[idx].cumsum()[idx.argsort()]

设置输入和计时 -

In [212]: # Inputs
     ...: N = 5000
     ...: A = np.random.rand(N)
     ...: B = np.random.rand(N)
     ...: 

In [213]: %timeit org_app(A,B)
     ...: %timeit sum_based(A,B)
     ...: %timeit dot_based(A,B)
     ...: %timeit bincount_based(A,B)
     ...: %timeit org_app_modified(A,B)
     ...: %timeit cumsum_trick(A,B)
     ...: 
1 loops, best of 3: 266 ms per loop
1 loops, best of 3: 411 ms per loop
1 loops, best of 3: 322 ms per loop
1 loops, best of 3: 1.01 s per loop
10 loops, best of 3: 196 ms per loop
1000 loops, best of 3: 835 µs per loop