无需遍历所有索引填充numpy数组：替代嵌套循环计算矩阵加权和范数

完全可以通过numpy的广播向量化机制替代嵌套循环，完全避免Python层的索引遍历，运算效率会有量级提升。

实现思路 #

• 给四个权重系数i/j/p/q分别构造适配广播维度的数组，确保权重和对应矩阵相乘后可以自动拓展到匹配高维输出的形状
• 直接做矩阵加权和的广播运算，得到形状为(8,7,6,5,3,3)的高维数组
• 调用np.linalg.norm时指定在最后两个矩阵维度上计算范数，直接得到目标形状的结果

示例代码 #

import numpy as np

# 原始输入矩阵
mat1, mat2, mat3, mat4 = np.random.rand(3, 3), np.random.rand(3, 3), np.random.rand(3, 3), np.random.rand(3, 3)

# 构造带广播维度的权重数组
i = np.arange(8).reshape(8, 1, 1, 1, 1, 1)
j = np.arange(7).reshape(1, 7, 1, 1, 1, 1)
p = np.arange(6).reshape(1, 1, 6, 1, 1, 1)
q = np.arange(5).reshape(1, 1, 1, 5, 1, 1)

# 向量化计算加权和 + 范数
res_vectorized = np.linalg.norm(i * mat1 + j * mat2 + p * mat3 + q * mat4, axis=(-2, -1))

正确性验证 #

你可以用下面的代码和原循环生成的结果做对比，二者输出完全一致：

# 原始循环实现
res_loop = np.empty((8, 7, 6, 5))
for i in range(8):
    for j in range(7):
        for p in range(6):
            for q in range(5):
                 res_loop[i, j, p, q] = np.linalg.norm(i * mat1 + j * mat2 + p * mat3 + q * mat4)

# 验证结果一致
print(np.allclose(res_loop, res_vectorized))  # 输出为True

如果后续使用时维度规模变大，需要注意向量化版本的内存占用，常规规模下其运算速度远快于Python层嵌套循环。

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