NumPy三维数组如何实现相邻子矩阵分组逐元素求和

你的当前代码存在分组配对错位的问题，无法得到预期输出，具体问题和正确方案如下：

原代码的逻辑偏差 #

你原本的设计思路是对切片a[:2]内部的两个子矩阵求和、对切片a[2:]内部的两个子矩阵求和，但当前代码把两个切片作为整体做了逐元素相加：

• a[:2] 是最外层维度索引0、1的两个子矩阵，shape为(2,2,2)
• a[2:] 是最外层维度索引2、3的两个子矩阵，shape为(2,2,2)

两者逐元素相加时，是索引0的子矩阵和索引2的子矩阵对应位置相加、索引1的子矩阵和索引3的子矩阵对应位置相加，最终输出为：

array([[[10, 12],
        [14, 16]],

       [[18, 20],
        [22, 24]]])

和你期望的输出存在明显偏差。

符合需求的正确实现 #

你的核心需求是：将最外层维度的4个矩阵按「相邻两个为一组」的规则，每组内的两个矩阵对应位置求和，最终得到2组求和结果，shape为(2,2,2)。有两种常用实现方式：

直观索引计算（逻辑最清晰） #

直接按分组取对应矩阵做求和，再堆叠为最终数组：

# 第一组：索引0和索引1的子矩阵求和
group1 = a[0] + a[1]
# 第二组：索引2和索引3的子矩阵求和
group2 = a[2] + a[3]
# 堆叠为目标shape的数组
output = np.stack([group1, group2])

通用维度变换（适合批量分组场景） #

如果分组数量多、不想手写索引，可以通过维度重排+求和的方式实现，本质是把原数组的第一维拆分为「分组数、组内元素数」两个维度，再在组内维度上做聚合求和：

# 原a shape为(4,2,2)，拆分为(2个分组, 每组2个矩阵, 矩阵行, 矩阵列)的结构
# 对轴1（也就是组内的矩阵维度）做求和，即可得到每组的求和结果
output = a.reshape(2, 2, 2, 2).sum(axis=1)

两种方式最终输出的结果均和你给出的预期值完全一致：

array([[[ 6,  8],
        [10, 12]],

       [[22, 24],
        [26, 28]]])

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