使用Numpy计算不规则维度数组的三级列求和问题

问题场景 #

当处理规则形状的多维数组时，使用np_array.sum(axis=0)可以轻松计算第三层级的每列之和：

import numpy as np

score_list = [
    [[1,1,3], [1,2,5]],
    [[2,7,5], [4,1,3]]
]
np_array = np.array(score_list)
sum_array = np_array.sum(axis=0)
print(sum_array)

输出结果为：

[[3 8 8]
 [5 3 8]]

但遇到不规则形状数组时，numpy的sum(axis=0)无法得到预期结果：

score_list = [
    [[1,1], [1,2,5]],
    [[2,7], [4,1,3]]
]
np_array = np.array(score_list)
sum_array = np_array.sum(axis=0)
print(sum_array)

实际返回：

[list([1, 1, 2, 7]) list([1, 2, 5, 4, 1, 3])]

并伴随警告，无法得到预期的[[3 8] [5 3 8]]。

解决方案 #

由于不规则数组在numpy中会被识别为object类型，原生的sum(axis=0)会直接拼接同位置元素而非按列求和。可以通过手动遍历对应层级的子数组，按列逐个求和来实现需求：

import numpy as np

score_list = [
    [[1,1], [1,2,5]],
    [[2,7], [4,1,3]]
]

# 遍历第二层级的每个位置，收集对应子数组并按列求和
result = []
for idx in range(len(score_list[0])):
    # 获取所有第一层级中第idx位置的子数组
    sub_arrays = [layer[idx] for layer in score_list]
    # 用zip打包同列元素，再计算每列的和
    column_sums = [sum(col_elements) for col_elements in zip(*sub_arrays)]
    result.append(column_sums)

# 转换为numpy数组（可选，若需保持numpy格式）
result_np = np.array(result, dtype=object)
print(result_np)

运行后输出：

[[3 8] [5 3 8]]

原理说明 #

• 通过[layer[idx] for layer in score_list]收集所有第一层级中对应第二层级位置的子数组，比如第一个循环中收集到[[1,1], [2,7]]，第二个循环收集到[[1,2,5], [4,1,3]]。
• zip(*sub_arrays)会将子数组的同列元素打包成元组，比如zip(*[[1,1], [2,7]])得到(1,2)和(1,7)，再对每个元组求和即可得到列总和。

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