np.array 的运行总和

Question

我需要计算 np.array 的运行总和。假设

a

是给定的长度为

l

的 np.array，并且

n

是一个整数步长，那么我需要一个长度为

b

的数组

l - n + 1

，其中

b= [a[0]+a[1]+...+a[n-1], a[1]+a[2]+...+a[n], ..., a[ l-n]+a[l-n+1]+...+a[l-1]]

我尝试了以下方法：

def run_sum(a, n):     
    a_cumsum = np.r_[0, np.cumsum(a)]     
    b = a_cumsum[n:] - a_cumsum[:-n]     
    return b

它工作得很好，但是如果数组

a

有非常不同的元素（例如其中一些是

1e8

，其中一些是

1e-8

），就像我的情况一样，那么这个方法会给出太多的错误，这些错误会累积在

中np.cumsum

。您对如何直接执行此操作而不使用

np.cumsum

有什么想法吗？

由于

a

有长度

int(1e8)

我不想使用

for

循环。想要类似的东西：

ind = np.arange(a.size - n + 1)
b[ind] = np.sum(a[ind:ind+n])

但它不起作用（

Answer 1

如果你想做类似 for 循环的事情，你可以使用

sliding_window_view

:

from numpy.lib.stride_tricks import sliding_window_view as swv

b = swv(a, n).sum(axis=1)

然而，这给出了与您的函数非常相似的结果：

# generate an input with a large variability
np.random.seed(0)
a = 1/np.random.random(size=100_000_000)*1e-8

a.min(), a.max()
# (1.0000000074212765e-08, 0.9650719974971247)

b1 = run_sum(a, n)
b2 = swv(a, n).sum(axis=1)

# are the two outputs almost equal?
np.allclose(b1, b2)
# True

# what is the greatest difference?
np.abs(b2-b1).max()
# 7.059628683189656e-15

但速度相当慢：

# run_sum(a, n)
895 ms ± 11.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# swv(a, n).sum(axis=1)
2.29 s ± 93.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)