如何利用reshape将Numpy二维数组分步求和结果转为二维数组？

原始代码与输出 #

原始实现代码：

import numpy as np

a = np.array([[5, 1, 8, 1, 6, 1, 3, 2],[2, 3, 4, 1, 6, 1, 4, 2]])
n = 2
[(a[0:2, i:i+n]).sum(axis=1) for i in range(0,a.shape[1],n)]

输出结果为一个包含多个子数组的列表：

[array([6, 5]), array([9, 5]), array([7, 7]), array([5, 6])]

需求：将结果转换为单一的二维数组，并用更优雅的方式通过reshape实现。

优雅实现方案 #

利用reshape将数组按指定列数分块，再求和并调整维度，直接得到目标二维数组：

import numpy as np

a = np.array([[5, 1, 8, 1, 6, 1, 3, 2],[2, 3, 4, 1, 6, 1, 4, 2]])
n = 2

# 一步到位实现分块求和并生成二维数组
result = a.reshape(a.shape[0], n, -1).sum(axis=1).T

print(result)

运行输出：

array([[6, 5],
       [9, 5],
       [7, 7],
       [5, 6]])

代码解析 #

1. reshape分块：a.reshape(a.shape[0], n, -1)将原数组（shape为(2,8)）重塑为三维数组(2,2,4)，其中：
   • 第一维保留原数组的行数（2行）
   • 第二维是每块的列数（n=2）
   • 第三维由NumPy自动计算分块数量（8//2=4）
2. 分块求和：sum(axis=1)对每个分块内的元素求和，消去分块维度，得到shape为(2,4)的数组
3. 转置调整结构：.T将数组转置为(4,2)，与原列表中子数组的排列结构完全匹配

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