Different length sum on numpy array axis

user_4401979 0
Viewed 43

Suppose to have, a numpy 3D tensor D of dimension r x c x d, such as:

r = 2
c = 3
d = 3    
D = np.array([[[1, 5, 3], [1, 2, 5], [1, 4, 3]], [[1, 1, 6], [3, 1, 7], [5, 1, 3]]])

array([[[1, 5, 3],
        [1, 2, 5],
        [1, 4, 3]],

       [[1, 1, 6],
        [3, 1, 7],
        [5, 1, 3]]])

and a 2D integer matrix Q of dimensions r x c, such as:

Q = np.array([[1, 1, 2], [2, 1, 2]])

array([[1, 1, 2],
       [2, 1, 2]])

where every element in Q is less than d.

I need to sum the first Q[r_i][c_i] element of the third dimension of matrix D for every 0 < r_i < r and 0 < c_i < c. The expected results (Res) using the example above is a 2D matrix of r x c (2x3):

Res = np.array([[6, 3, 8], [8, 4, 5]])
array([[6, 3, 8],
       [8, 4, 5]])

My actual solution is using a list comprehension looping over r_i and c_i:

r = 2
c = 3
res = np.array([[np.sum(D[r_i, c_i, :Q[r_i, c_i]+1]) for c_i in range(c)] for r_i in range(r)])

There is a more efficient or elegant solution to solve this problem?

numpy python
1 Answers

Let us try:

# this is equivalent to double loop on r_i, c_i
x,y = np.ogrid[:r, :c]

# we take the cumsum on the last axis, 
# then extract the Q[r_i, c_i]'th sum at r_i, c_i
out = D.cumsum(axis=-1)[x,y, Q]

Output:

array([[6, 3, 8],
       [8, 4, 9]])

Cross check

np.allclose(out, res)
# True
user_42384080
Related