Суммирование переменных в массиве Python

Question

Я ищу решение для суммирования на столбец в двумерном массиве («a» в примере ниже) и начиная с позиции ячейки, как определено в другом массиве 1D («ref» в примере ниже).

Я пробовал следующее:

import numpy as np

a = np.arange(20).reshape(5, 4)
print(a)                         # representing an original large 2D array
ref = np.array([0, 2, 4, 1])     # reference array for defining start of sum
s = a.sum(axis=0)
print(s)    # Works: sums all elements per column
s = a[2:].sum(axis=0)
print(s)    # Works as well: sum from the third element till end per column

# This is what I look for: sum per column starting at element defined by ref[]
s = np.zeros(4).astype(int)      # makes an empty 1D array
for i in np.arange(4):           # for each column
    for j in np.arange(ref[i], 5):
        s[i] += a[j, i]          # sums all elements from ref till end (i.e. 5)

print(s)    # This is the desired outcome

for i in np.arange(4):
    s = a[ref[i]:].sum(axis=0)

print(s)    # No good; same as a[ref[4]:].sum(axis=0) and here ref[4] = 1

s = np.zeros(4).astype(int)      # makes an empty 1D array
for i in np.arange(4):
    s[i] = np.sum(a[ref[i]:, i])

print(s)    # Yes; this is also the desired outcome

Возможно ли реализовать это без использования цикла for? Есть ли у numpy функции для выполнения этого за один шаг?

s = a[ref:].sum(axis=0)

Это было бы неплохо, но не работает.

Спасибо за ваше время!

Answer 1

Базовое решение на основе np.cumsum:

In [1]: a = np.arange(15).reshape(5, 3)

In [2]: res = np.array([0, 2, 3])

In [3]: b = np.cumsum(a, axis=0)

In [4]: b
Out[4]: 
array([[ 0,  1,  2],
       [ 3,  5,  7],
       [ 9, 12, 15],
       [18, 22, 26],
       [30, 35, 40]])

In [5]: a
Out[5]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])


In [6]: b[res, np.arange(a.shape[1])]
Out[6]: array([ 0, 12, 26])

In [7]: b[-1, :] - b[res, np.arange(a.shape[1])]
Out[7]: array([30, 23, 14])

, поэтому он не дает нам желаемого результата: нам нужно добавить первую строку нулей в b:

In [13]: b = np.vstack([np.zeros((1, a.shape[1])), b])

In [14]: b
Out[14]: 
array([[  0.,   0.,   0.],
       [  0.,   1.,   2.],
       [  3.,   5.,   7.],
       [  9.,  12.,  15.],
       [ 18.,  22.,  26.],
       [ 30.,  35.,  40.]])

In [17]: b[-1, :] - b[res, np.arange(a.shape[1])]
Out[17]: array([ 30.,  30.,  25.])

что, я полагаю, является желаемым результатом.