在NumPy中进行多维累加求和

Question

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

2D: xᶦʲ = ∑xᵐⁿ，其中m≥i且n≥j

3D: xᶦʲᵏ = ∑xᵐⁿᵒ，其中m≥i且n≥j且o≥k

2D示例：

1 1 0       2  1  0
1 1 1  ->   5  3  1
1 1 1       8  5  2

1 2 3       6  5  3
4 5 6  ->  21 16  9
7 8 9      45 33 18

在三维中的例子：

1 1 1       3   2   1
1 1 1       6   4   2
1 1 1       9   6   3

1 1 1       6   4   2
1 1 1  ->  12   8   4
1 1 1      18  12   6

1 1 1       9   6   3
1 1 1      18  12   6
1 1 1      27  18   9

- user2685230

Answer 1

对于那些希望获得类似于“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（感谢他！）

- ch271828n

---

完美的，无价的！ - jtlz2

Answer 2

沿最后一个轴翻转，沿相同轴累加和，再将其翻转回来，最后沿第二个到倒数第二个轴开始累加和，直到第一个轴 -

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]]])

- Divakar

---

当然，谢谢！对于我的用例，完整的4D解决方案是：a[..., ::-1].cumsum(-1)[..., ::-1].cumsum(-2).cumsum(-3).cumsum(-4)。 - user2685230

@user2685230 编辑了一个基于相同思路的通用解决方案。 - Divakar

这太神奇了，真应该称之为魔方:D - jtlz2

Answer 3

这里有一个通用解决方案。我按照描述，而不是示例进行操作，即垂直显示的顺序是自上而下而不是自下而上。

import itertools as it
import functools as ft

ft.reduce(np.cumsum, it.chain((a[a.ndim*(np.s_[::-1],)],), range(a.ndim)))[a.ndim*(np.s_[::-1],)]

或者原地修改：

for i in range(a.ndim):
    b = a.swapaxes(0, i)[::-1]
    b.cumsum(axis=0, out=b)

- Paul Panzer

---

你的原地版本似乎有一个 bug？在后续的迭代中 b 将被覆盖。 - shaunc

@shaunc "[...] bug? bis overwritten [...]" ??? 这不就是 in-place 的意思吗？ - Paul Panzer

我的意思是——假设a.ndim > 1。在每次迭代中，b不依赖于先前计算出的b的值，而是从a获得一个新的副本，忽略了以前的工作。 - shaunc