使用递归简化嵌套循环：高维数组均值计算实现咨询

我们需要对任意高维数组，计算前两维构成的二维矩阵的均值，将结果存入对应剩余维度位置的输出数组中。当数组维度提升至100维时，嵌套循环会异常繁琐，以下是基于递归的实现方案：

核心思路 #

• 递归终止条件：当输入数组为二维矩阵时，直接返回该矩阵的均值。
• 递归递进逻辑：对数组从第3维开始的每个维度，遍历所有索引，对每个索引对应的子数组递归执行均值计算，最终将结果汇总到输出数组中。

代码实现 #

import numpy as np

def recursive_mean(x):
    # 终止条件：当前是二维矩阵，直接返回均值
    if x.ndim == 2:
        return np.mean(x)
    
    # 确定剩余维度的形状，初始化输出数组
    remaining_dims_shape = x.shape[2:]
    out = np.zeros(remaining_dims_shape)
    
    # 遍历当前第一个剩余维度的所有索引
    for idx in range(x.shape[2]):
        # 递归处理子数组，并将结果赋值到输出对应位置
        if len(remaining_dims_shape) == 1:
            out[idx] = recursive_mean(x[:, :, idx])
        else:
            out[idx, ...] = recursive_mean(x[:, :, idx, ...])
    
    return out

# 测试示例
# 3维数组验证
x = np.random.normal(0, 1, [4, 4, 3])
out = recursive_mean(x)
print("3维数组输出形状:", out.shape)
print("与循环结果一致性:", np.allclose(out, np.array([np.mean(x[:, :, k]) for k in range(3)])))

# 4维数组验证
x2 = np.random.normal(0, 1, [5, 5, 4, 6])
out2 = recursive_mean(x2)
print("4维数组输出形状:", out2.shape)
# 循环实现的结果用于对比
out2_loop = np.zeros(x2.shape[2:])
for k in range(4):
    for l in range(6):
        out2_loop[k, l] = np.mean(x2[:, :, k, l])
print("与循环结果一致性:", np.allclose(out2, out2_loop))

# 高维数组测试（如5维）
x5 = np.random.normal(0, 1, [2, 2, 3, 4, 5])
out5 = recursive_mean(x5)
print("5维数组输出形状:", out5.shape)

代码说明 #

• 函数recursive_mean自动适配任意高于二维的数组，无需手动编写多层嵌套循环。
• 使用...（Ellipsis）处理多维索引，确保不管剩余维度数量多少，都能正确赋值到输出数组的对应位置。
• 每一层递归都会将数组维度减少1，直到触发二维矩阵的终止条件，完美适配100维等高维场景。

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