使用递归简化嵌套循环

Question

假设我们有一个三维数组

x

。对于每个

k

，我想将

x[:, :, k]

的平均值存储为

out[k]

。如果是 3D 或 4D 矩阵，任务很简单，for 循环如下：

x = np.random.normal(0, 1, [4, 4, 3])
out = np.zeros(x[0, 0, :].shape)
for k in range(3):
    out[k] =  np.mean(x[:, :, k])
    
x2 = np.random.normal(0, 1, [5, 5, 4, 6])
out2 = np.zeros(x2[0, 0, :, :].shape)
for k in range(4):
    for l in range(6):
        out2[k, l] = np.mean(x2[:, :, k, l])

但是如果我想获得更高的维度（假设我想覆盖多达 100 个维度），代码会变得很难看，因为嵌套循环的数量增加了。

我知道我必须递归地使用二维矩阵作为我的基本情况，但我不知道如何实现它。

Answer 1

你似乎想取前两个轴的平均值。你可以简单地通过

x.mean(axis=(0, 1))

来做到这一点

例子：

• out*

  是你的方法计算出来的平均值

• out*_np

  是我上面显示的

  np.mean

  计算的平均值。

x = np.random.normal(0, 1, [4, 4, 3])
out = np.zeros(x[0, 0, :].shape)
out_np = x.mean(axis=(0,1))
for k in range(3):
    out[k] =  np.mean(x[:, :, k])

print(np.allclose(out, out_np)) # True

x2 = np.random.normal(0, 1, [5, 5, 4, 6])
out2 = np.zeros(x2[0, 0, :, :].shape)
out2_np = x2.mean(axis=(0,1))
for k in range(4):
    for l in range(6):
        out2[k, l] = np.mean(x2[:, :, k, l])

print(np.allclose(out2, out2_np)) # True

---

要将此推广到在 2D 矩阵上运行的任何函数，您可以将原始矩阵重塑为 3D 矩阵，在最后一个轴上应用您的方法，并将结果重塑为所需的形状：

def apply_2dfunc(X, func):
    X_r = X.reshape((*X.shape[0:2], -1))
    res = np.zeros((X_r.shape[-1],))
    for i in range(len(res)):
        res[i] = func(X_r[:, :, i])

    return res.reshape(X.shape[2:])

以你的

func=np.mean

为例，

out_g = apply_2dfunc(x, np.mean)
print(np.allclose(out_g, out)) # True


out2_g = apply_2dfunc(x2, np.mean)
print(np.allclose(out2_g, out2)) # True

Answer 2

这种情况不需要递归。原来我之前的答案非常接近，只需要一些转置。

x2 = np.random.normal(0, 1, [2, 2, 7, 5, 5, 4, 6])

dims = len(x2.shape)

out = np.mean(x2.T, axis = tuple(range(-2, -1))).T


print(out[0,0,0,0,0], np.mean(x2[:, :, 0, 0, 0, 0, 0]))

值得注意的是，

out

和

np.mean

之间的值有时在最后三位数字上略有偏差，因此不知何故会发生一些奇怪的精度损失。

如果你必须递归地做...

x2 = np.random.normal(0, 1, [2, 2, 7, 5, 5, 4, 6])

def recursive_mean(x, out = None):
    if out is None:
        out = np.zeros(x.shape[2:])
        
    if len(x.shape) > 3:
        for i in range(x.shape[2]):
            out[i] = recursive_mean(x[:, :, i], out[i])
            
    else:
        for i in range(x.shape[2]):
            out[i] = np.mean(x[:, :, i])
        
        
    return out

out2 = recursive_mean(x2)

print(out2[0,0,0,0,0], np.mean(x2[:, :, 0, 0, 0, 0, 0]))
print(out2[3, 1, 2, 0, 4], np.mean(x2[:, :, 3, 1, 2, 0, 4]))