Multidimensional cumulative sum in numpy

问题

I want to be able to calculate the cumulative sum of a large n-dimensional numpy array. The value of each element in the final array should be the sum of all elements which have indices greater than or equal to the current element.

2D: xᶦʲ = ∑xᵐⁿ ∀ m ≥ i and n ≥ j

3D: xᶦʲᵏ = ∑xᵐⁿᵒ ∀ m ≥ i and n ≥ j and o ≥ k

Examples in 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

Example in 3D:

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

回答1:

Flip along the last axis, cumsum along the same, flip it back and finally cumsum along the second last axis onwards until the first axis -

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

Sample 2D case run -

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

Sample 3D case run -

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

回答2:

For those who want a "numpy-like" cumsum where the top-left corner is smallest:

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

Modified from @Divakar (thanks to him!)

回答3:

Here is a general solution. I'm going by the description, not the examples, i.e. order of vertical display is top down not bottom up:

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

Or in-place:

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

来源：https://stackoverflow.com/questions/50463366/multidimensional-cumulative-sum-in-numpy

标签

python

arrays

numpy