如何通过for循环实现三维数组的指定规则累积求和

原代码问题分析 #

• 求和时未正确替换循环下标：内层求和的jj/kk循环中，仍然使用外层的固定j/k下标获取x的值，导致求和范围错误
• 循环顺序不符合计算逻辑：应该优先遍历独立的i维度，再按k从低到高、j从小到大的顺序计算，利用已计算的前缀和简化运算
• 未匹配给定的累加规则：k=1的计算需要先加完k=0层所有j的总和，再加k=1层到当前j的前缀和，原代码没有区分两层k的计算逻辑

正确实现代码 #

# 初始化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]]]
# 初始化和x同维度的y数组，初始值为0
y = [[[0]*2 for _ in range(5)] for _ in range(4)]

# 按规则计算
for i in range(4):
    # 处理k=0层的前缀和
    prefix_k0 = 0
    for j in range(5):
        prefix_k0 += x[i][j][0]
        y[i][j][0] = prefix_k0
    # 提前计算k=0层所有j的总和，用于k=1层计算
    sum_k0_all = prefix_k0
    # 处理k=1层的前缀和
    prefix_k1 = 0
    for j in range(5):
        prefix_k1 += x[i][j][1]
        y[i][j][1] = sum_k0_all + prefix_k1

运行结果验证 #

运行上述代码后输出的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]]]

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