Python：多大型矩阵逐元素平均的高效实现方法问询

当然可以直接用numpy的向量化操作替代低效的循环，这也是numpy处理大规模数组的核心优势。

核心思路 #

你的rho_tm数组形状应该是(N, N, T, M)（其中N=1000，T=500，M=10000），你需要对**最后一个维度（M维度）**做逐元素平均。numpy的np.mean支持通过axis参数指定要计算平均的维度，完全不需要手动遍历i和j。

替代代码 #

直接用一行代码替换你原来的三重循环：

rho_t = np.mean(rho_tm, axis=-1, dtype=np.complex64)

为什么这更快？ #

• 原来的代码嵌套了三层循环，其中i和j的循环要执行1000×1000=100万次，Python的循环本身速度很慢，处理这么大的次数会非常耗时。
• numpy的内置函数是底层C语言实现的向量化操作，能一次性对整个数组的对应维度做计算，避免了Python循环的开销，效率会提升几个数量级。

验证形状 #

执行完上面的代码后，rho_t的形状就是(N, N, T)，和你原来初始化的rho_t形状完全一致，每个元素都是对应位置上M个矩阵的平均值，和手动循环计算的结果完全相同。

内容的提问来源于stack exchange，提问作者J.Agusti