如何将求和向量堆叠至NumPy三维数组？两类场景实现问询

场景1：将axis=1的求和结果作为内部组的第四个向量堆叠 #

原三维数组g的形状为(2, 3, 5)，沿第1轴求和得到的gs形状是(2, 5)。要将其作为每个内部组的第四个向量，需要先给gs增加一个中间维度，使其形状匹配为(2, 1, 5)，再与原数组在axis=1维度拼接：

import numpy as np

# 定义原数组
g = np.array([[[ 1.,  1.,  1.,  1.,  1.],
               [ 0.,  0.,  0.,  0.,  0.],
               [ 1.,  2.,  3.,  4.,  6.]],

              [[ 0.,  0.,  0.,  0.,  0.],
               [11., 22., 33., 44., 66.],
               [ 0.,  0.,  0.,  0.,  0.]]])

# 沿axis=1计算求和
gs = g.sum(axis=1)
# 扩展维度，适配拼接要求
gs_expanded = gs[:, np.newaxis, :]
# 在axis=1维度完成堆叠
result1 = np.concatenate([g, gs_expanded], axis=1)

print(result1)

执行后将得到预期的(2, 4, 5)形状数组，每个内部组新增了对应求和向量。

场景2：将axis=0的求和结果作为新的外部组堆叠 #

沿第0轴求和得到的gss形状是(3, 5)，要将其作为新的外部组加入原数组，需先给gss增加第一个维度，使其形状变为(1, 3, 5)，再与原数组在axis=0维度拼接：

# 沿axis=0计算求和
gss = g.sum(axis=0)
# 扩展维度，适配拼接要求
gss_expanded = gss[np.newaxis, :, :]
# 在axis=0维度完成堆叠
result2 = np.concatenate([g, gss_expanded], axis=0)

print(result2)

执行后将得到预期的(3, 3, 5)形状数组，新增了由原数组沿第0轴求和得到的外部组。

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