numpy中的多维累积和

Question

我希望能够计算一个大的 n 维 numpy 数组的累积和。最终数组中每个元素的值应该是索引大于或等于当前元素的所有元素的总和。

\n\n

2D: x\xe1\xb6\xa6\xca\xb2 = \xe2\x88\x91x\xe1\xb5\x90\xe2\x81\xbf \xe2\x88\x80 m \xe2\x89\xa5 i 和 n \xe2 \x89\xa5 j

\n\n

3D: x\xe1\xb6\xa6\xca\xb2\xe1\xb5\x8f = \xe2\x88\x91x\xe1\xb5\x90\xe2\x81\xbf\xe1\xb5\x92 \xe2\x88\x80 m \xe2\x89\xa5 i 和 n \xe2\x89\xa5 j 和 o \xe2\x89\xa5 k

\n\n

二维示例：

\n\n

1 1 0       2  1  0\n1 1 1  ->   5  3  1\n1 1 1       8  5  2\n\n1 2 3       6  5  3\n4 5 6  ->  21 16  9\n7 8 9      45 33 18\n

\n\n

3D 示例：

\n\n

1 1 1       3   2   1\n1 1 1       6   4   2\n1 1 1       9   6   3\n\n1 1 1       6   4   2\n1 1 1  ->  12   8   4\n1 1 1      18  12   6\n\n1 1 1       9   6   3\n1 1 1      18  12   6\n1 1 1      27  18   9\n

\n

Answer 1

沿最后一个轴翻转，沿同一条轴进行 cumsum，再翻转回来，最后沿倒数第二个轴向前进行 cumsum，直到第一个轴 -

def multidim_cumsum(a):
    out = a[...,::-1].cumsum(-1)[...,::-1]
    for i in range(2,a.ndim+1):
        np.cumsum(out, axis=-i, out=out)
    return out

示例 2D 案例运行 -

In [107]: a
Out[107]: 
array([[1, 1, 0],
       [1, 1, 1],
       [1, 1, 1]])

In [108]: multidim_cumsum(a)
Out[108]: 
array([[2, 1, 0],
       [5, 3, 1],
       [8, 5, 2]])

示例 3D 案例运行 -

In [110]: a
Out[110]: 
array([[[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]],

       [[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]],

       [[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]]])

In [111]: multidim_cumsum(a)
Out[111]: 
array([[[ 3,  2,  1],
        [ 6,  4,  2],
        [ 9,  6,  3]],

       [[ 6,  4,  2],
        [12,  8,  4],
        [18, 12,  6]],

       [[ 9,  6,  3],
        [18, 12,  6],
        [27, 18,  9]]])

Answer 2

对于那些想要numpy左上角最小的“类似”累积和的人：

def multidim_cumsum(a):
    out = a.cumsum(-1)
    for i in range(2,a.ndim+1):
        np.cumsum(out, axis=-i, out=out)
    return out

修改自@Divakar（感谢他！）