CODEWITHSUNDEEP
pythonarraysnumpy
user2685230
user2685230

Reputation: 107

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

Upvotes: 6

Views: 3710

Answers (3)

ch271828n
ch271828n

Reputation: 17607

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

Upvotes: 5

Paul Panzer
Paul Panzer

Reputation: 53069

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)

Upvotes: 2

Divakar
Divakar

Reputation: 221614

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

Upvotes: 5

Related Questions