How to sum an ndarray over ranges bounded by other indexes-CodePudding

For an array of multiple dimensions, I would like to sum along some dimensions, with the sum range defined by other dimension indexes. Here is an example:

>>> import numpy as np
>>> x = np.arange(2*3*4).reshape((2,3,4))
>>> x
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])
>>> wanted = [[sum(x[i,j,i:j]) for j in range(x.shape[1])] for i in range(x.shape[0])]
>>> wanted
[[0, 4, 17], [0, 0, 21]]

Is there a more efficient way to do it without for loops or list comprehension? My array is quite large.

CodePudding user response：

You can use boolean masks:

# get lower triangles
m1 = np.arange(x.shape[1])[:,None]>np.arange(x.shape[2])

# get columns index >= depth index
m2 = np.arange(x.shape[2])>=np.arange(x.shape[0])[:,None,None]

# combine both mask to form 3D mask
mask = m1 & m2

out = np.where(mask, x, 0).sum(axis=2)

output:

array([[ 0,  4, 17],
       [ 0,  0, 21]])

Masks:

# m1
array([[False, False, False, False],
       [ True, False, False, False],
       [ True,  True, False, False]])

# m2
array([[[ True,  True,  True,  True]],

       [[False,  True,  True,  True]]])

# mask
array([[[False, False, False, False],
        [ True, False, False, False],
        [ True,  True, False, False]],

       [[False, False, False, False],
        [False, False, False, False],
        [False,  True, False, False]]])