Numpy计算3D矩阵指定轴平均值 输出(P,T)形状2D矩阵

需求说明 #

现有3D NumPy矩阵stocks，形状为(A, P, T)，各维度含义：

• A：股票代码总数量
• P：单一时点下的股票价格维度数量
• T：时间维度
  需要返回形状为(P, T)的2D矩阵，每个元素为对应时点、对应价格位置上所有股票价格的平均值。
  示例矩阵定义代码如下：

stocks = np.array([
    [ [1,2,3],[4,5,6],[7,8,9] ],
    [ [10,11,12],[13,14,15],[16,17,18] ],
    [ [19,20,21],[22,23,24],[25,26,27] ],
    [ [28,29,30],[31,32,33],[34,35,36] ]
])

具体实现 #

直接调用NumPy数组的mean方法，指定沿着第0轴（即对应股票数量的A维度，NumPy维度索引从0开始计数）计算平均值即可。计算完成后会自动消去第0轴，保留(P, T)的维度形状，完全匹配输出要求。
完整实现代码：

import numpy as np

stocks = np.array([
    [ [1,2,3],[4,5,6],[7,8,9] ],
    [ [10,11,12],[13,14,15],[16,17,18] ],
    [ [19,20,21],[22,23,24],[25,26,27] ],
    [ [28,29,30],[31,32,33],[34,35,36] ]
])

# 沿股票维度A（axis=0）计算所有股票的价格均值
avg_prices = stocks.mean(axis=0)

结果验证 #

针对上述示例矩阵，计算得到的avg_prices形状为(3,3)，数值如下，和手动计算结果完全一致：

array([[14.5, 15.5, 16.5],
       [17.5, 18.5, 19.5],
       [20.5, 21.5, 22.5]])

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