Python中的可变部分数组求和

时间:2018-05-03 19:40:05

标签: python arrays numpy sum

我正在寻找一种解决方案,在2D数组中对每列求和(在下面的示例中为“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)

这很好,但不起作用。

感谢您的时间!

1 个答案:

答案 0 :(得分: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.])

我认为,这是所需的输出。