Python如何仅用for循环与sum函数实现三维数组指定规则累积求和

我有一个三维数组x，结构如下：

x=
[[[1, 6], [2, 7], [3, 8], [4, 9], [5, 10]],
 [[1, 6], [2, 7], [3, 8], [4, 9], [5, 10]],
 [[1, 6], [2, 7], [3, 8], [4, 9], [5, 10]],
 [[1, 6], [2, 7], [3, 8], [4, 9], [5, 10]]]

需要得到如下的累积求和结果y：

y=
[[[1, 21], [3, 28], [6, 36], [10, 45], [15, 55]],
 [[1, 21], [3, 28], [6, 36], [10, 45], [15, 55]],
 [[1, 21], [3, 28], [6, 36], [10, 45], [15, 55]],
 [[1, 21], [3, 28], [6, 36], [10, 45], [15, 55]]]

实现要求 #

• 仅使用for循环和sum内置函数
• 不能使用numpy、itertools等专用函数
• 不允许使用额外临时变量

已尝试代码与问题 #

我编写的代码如下：

for k in range(0,1):
    for j in range(0,5):
        for i in range(0,4):
            y[i][j][k]=sum(sum(x[i][j][k] for jj in range(0,5) if jj<=j)for kk in range(0,1) if kk<=k)

运行后得到的结果不符合预期，错误结果如下：

y[i][j][k]=[[[1, 12], [3, 26], [6, 42], [10, 60], [15, 80]],
 [[1, 12], [3, 26], [6, 42], [10, 60], [15, 80]],
 [[1, 12], [3, 26], [6, 42], [10, 60], [15, 80]],
 [[1, 12], [3, 26], [6, 42], [10, 60], [15, 80]]]

预期计算规则 #

已知数组元素示例：

x[0][0][0]=1
x[0][1][0]=2
x[0][2][0]=3
x[0][3][0]=4
x[0][4][0]=5
x[0][0][1]=6
x[0][1][1]=7
x[0][2][1]=8
x[0][3][1]=9
x[0][4][1]=10

y的计算逻辑如下：

y[0][0][0]=x[0][0][0]=1
y[0][1][0]=x[0][0][0]+x[0][1][0]=3
y[0][2][0]=x[0][0][0]+x[0][1][0]+x[0][2][0]=6
y[0][3][0]=x[0][0][0]+x[0][1][0]+x[0][2][0]+x[0][3][0]=10
y[0][4][0]=x[0][0][0]+x[0][1][0]+x[0][2][0]+x[0][3][0]+x[0][4][0]=15

y[0][0][1]=x[0][0][0]+x[0][1][0]+x[0][2][0]+x[0][3][0]+x[0][4][0]+x[0][0][1]=21
y[0][1][1]=x[0][0][0]+x[0][1][0]+x[0][2][0]+x[0][3][0]+x[0][4][0]+x[0][0][1]+x[0][1][1]=28
y[0][2][1]=x[0][0][0]+x[0][1][0]+x[0][2][0]+x[0][3][0]+x[0][4][0]+x[0][0][1]+x[0][1][1]+x[0][2][1]=36
y[0][3][1]=x[0][0][0]+x[0][1][0]+x[0][2][0]+x[0][3][0]+x[0][4][0]+x[0][0][1]+x[0][1][1]+x[0][2][1]+x[0][3][1]=45
y[0][4][1]=x[0][0][0]+x[0][1][0]+x[0][2][0]+x[0][3][0]+x[0][4][0]+x[0][0][1]+x[0][1][1]+x[0][2][1]+x[0][3][1]+x[0][4][1]=55

请问如何编写符合要求的for循环代码实现该需求？

你原有代码的错误在于嵌套sum的取值逻辑不符合规则，没有区分k维度的不同计算要求。符合所有限制条件的实现代码如下：

# 先初始化与x结构一致的全0数组y
y = [[[0]*2 for _ in range(5)] for __ in range(4)]

for i in range(4):
    for j in range(5):
        # k=0维度：累加当前行前j+1个位置的k=0元素
        y[i][j][0] = sum(x[i][jj][0] for jj in range(j+1))
        # k=1维度：累加当前行所有k=0元素 + 前j+1个位置的k=1元素
        y[i][j][1] = sum(x[i][jj][0] for jj in range(5)) + sum(x[i][jj][1] for jj in range(j+1))

上述代码仅使用了for循环和sum内置函数，没有引入额外临时变量，也没有调用第三方库，运行后即可得到预期的y数组。

内容的提问来源于stack exchange，提问作者Deepan