高效的2d cumsum

Question

说我有这样的数组

>>> a = np.arange(1,8).reshape((1,-1))
>>> a
array([[1, 2, 3, 4, 5, 6, 7]])

并且我想为每个项目创建a一个"下4个项目的cumsum".也就是说,我的预期输出是

1,       2,      3, 4, 5, 6, 7, 8
1+2,     2+3,     ...
1+2+3    2+3+4    ...
1+2+3+4  2+3+4+5  ...

即包含的矩阵

1, 2, 3, 4, 5, 0, 0, 0
3, 5, 7, 9, 11,0, 0, 0
6, 9, 12,15,18,0, 0, 0
10,14,18,21,26,0, 0, 0

由于最后3个项目的cumsum操作无法正确完成,我期待0那里.我知道如何做一个单一的cumsum.实际上,阵列是

a[:4].cumsum().reshape((-1,1)); a[1:5].cumsum().reshape((-1,1))...

水平堆放.但是,我不知道如何以有效的方式做到这一点.这样做的好的矢量化numpy方式是什么？我也对scipy包装开放,只要它们numpy在效率或可读性方面占主导地位.

Answer 1

一种可能的方法是结合使用滚动窗口方法cumsum()。

例如：

from numpy.lib.stride_tricks import as_strided

a = np.arange(1, 9) # the starting array
slice_length = 4

然后你可以写：

arr = as_strided(a, (slice_length, len(a)), (a.strides[0], a.strides[0])).cumsum(axis=0)

这可以让您完成大部分工作，但要填写剩余的0值，您可以使用切片和分配来获得所需的输出：

arr[:, (1-slice_length):] = 0

然后你就有了数组：

>>> arr
array([[ 1,  2,  3,  4,  5,  0,  0,  0],
       [ 3,  5,  7,  9, 11,  0,  0,  0],
       [ 6,  9, 12, 15, 18,  0,  0,  0],
       [10, 14, 18, 22, 26,  0,  0,  0]])

我不知道是否有任何方法可以使用 NumPy 中的一种矢量化方法（即无需切片）来准确生成所需的输出。（accumulateat，有点像reduceat，添加到 NumPy 的 ufunc 中可能是一件有趣的事情......）