如何移除Python嵌套for循环优化numpy数组计算（不使用einsum）

直接利用numpy的轴变换（transpose）和向量化均值计算就能彻底移除外层循环，完全不需要einsum，核心思路是把所有全排列对应的数组变换一次性生成，再批量求均值。

具体实现步骤 #

3维数组的索引全排列一共6种，对应轴的排列组合是：[[0,1,2], [0,2,1], [1,0,2], [1,2,0], [2,0,1], [2,1,0]]。我们可以对每种排列用transpose调换数组的轴顺序，得到对应全排列索引的数组，最后堆叠起来求均值即可。

示例代码：

import numpy as np

# 假设t是你的3维输入数组，比如 t = np.random.rand(100, 100, 100)
# 定义所有3维索引的全排列对应的轴顺序
axis_perms = [[0,1,2], [0,2,1], [1,0,2], [1,2,0], [2,0,1], [2,1,0]]

# 生成所有排列后的数组
permuted_ts = [t.transpose(perm) for perm in axis_perms]
# 将这些数组沿新轴堆叠，再对该轴求均值
res = np.stack(permuted_ts, axis=0).mean(axis=0)

原理说明 #

• t.transpose([0,2,1]) 等价于遍历所有a,b,c取t[a,c,b]，这一步是numpy底层优化的向量化操作，比Python循环快几个数量级
• 堆叠后沿第0轴求均值，就是对每个位置[a,b,c]的6个全排列元素取平均，和你原来的循环/手动展开逻辑完全一致

这种方式完全没有外层循环，所有计算都由numpy的C底层实现，速度会比原循环版本提升非常明显。

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