如何计算矩阵截取后各列元素的求和？

问题描述 #

现有如下5×5对称矩阵：

Matrix = [[0., 15., 19., 18., 17.],
     [15., 0., 14., 12., 23.],
     [19., 14.,  0., 14., 21.],
     [18., 12., 14., 0., 14.],
     [17., 23., 21., 14.,  0.]]

通过以下代码截取矩阵后得到结果c：

c = [item[:index+1] for index, item in enumerate(Matrix[1:])]
# c的结果为：
# [array([15.]),
#  array([19., 14.]),
#  array([18., 12., 14.]),
#  array([17., 23., 21., 14.])]

需要计算c中对应列元素的求和，即：

• 第一列：15+19+18+17
• 第二列：14+12+23
• 第三列：14+21
• 第四列：14

解决方法 #

方法1：手动循环遍历列 #

直接按列遍历对应位置的元素累加，逻辑清晰直观：

sums = []
# 总共有4列需要计算
for col in range(4):
    total = 0.0
    # 第col列的元素存在于索引>=col的子数组中
    for i in range(col, len(c)):
        total += c[i][col]
    sums.append(total)

print(sums)  # 输出: [69.0, 49.0, 35.0, 14.0]

方法2：用zip补全后按列聚合 #

先将所有子数组补全到相同长度，再通过zip按列分组求和：

# 获取最长子数组的长度
max_len = len(c[-1])
# 用None补全每个子数组，统一长度
padded = [arr.tolist() + [None]*(max_len - len(arr)) for arr in c]
# 按列聚合，过滤掉补全的None后求和
sums = [sum(x for x in col if x is not None) for col in zip(*padded)]

print(sums)  # 输出: [69.0, 49.0, 35.0, 14.0]

方法3：利用numpy向量化操作（高效简洁） #

如果使用numpy数组，可借助其矩阵操作能力快速计算：

import numpy as np

max_len = len(c[-1])
# 创建空矩阵填充数据
np_c = np.zeros((len(c), max_len))
for i, arr in enumerate(c):
    np_c[i, :len(arr)] = arr

# 按列求和并转为列表
sums = np.sum(np_c, axis=0).tolist()

print(sums)  # 输出: [69.0, 49.0, 35.0, 14.0]

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