如何用Numpy内置函数替代矩阵乘法中的for循环优化计算

问题背景 #

给定形状为(m,n)的矩阵A，以及由数组lambd生成的对角矩阵lambda_diag，初始化代码如下：

import numpy as np
m, n = (44, 10)
A = np.random.random((m,n))
lambd = np.random.random(size=n)
lambda_diag = np.diag(lambd)

需要完成的计算逻辑：

• 对k从1到n，每次取A的前k列、lambd的前k个元素构造对角矩阵
• 计算矩阵乘积 A[:, :k] @ np.diag(lambd[:k]) @ A[:, :k].T 的对角线元素
• 将结果存入A_new的对应列（索引为k-1）

原实现采用for循环，效率较低且代码冗余：

A_new = np.empty_like(A)
A_new[:, 0] =  (A[:, 0] ** 2) * lambd[0]
for i in range(1, A.shape[1]):
    A_new[:, i] = np.diag(
        A[:, : i + 1]
        @ np.diag(lambd[: i + 1])
        @ A[:, : i + 1].T
    ).flatten()

需要用Numpy内置方案替代循环，提升效率。

优化方案 #

方法1：利用累积求和（最优解） #

观察矩阵乘积的对角线元素可以发现：A[:, :k] @ np.diag(lambd[:k]) @ A[:, :k].T的第i行对角线元素，本质是第i行前k个元素的平方与对应lambd值乘积的累加和。基于这个数学性质，可以直接构造贡献矩阵后做累积求和：

# 计算每个元素的贡献值：A[i,j]^2 * lambd[j]
contrib = A ** 2 * lambd[np.newaxis, :]
# 按列方向做累积求和，得到每一列对应的前k项累加结果
A_new = np.cumsum(contrib, axis=1)

该方法时间复杂度为O(mn)，远优于循环的O(mn²)，且代码极简。

方法2：使用np.einsum实现 #

如果偏好使用np.einsum，可以用它完成元素级的广播乘法，再配合累积求和：

# 用einsum计算每行元素平方与对应lambd的乘积
temp = np.einsum('ij,j->ij', A**2, lambd)
# 按列累积求和得到最终结果
A_new = np.cumsum(temp, axis=1)

此方法与方法1逻辑一致，仅在乘法步骤使用einsum实现，性能表现相近。

结果验证 #

可以通过以下代码验证优化方案与原循环结果的一致性：

# 原循环实现
A_new_loop = np.empty_like(A)
A_new_loop[:, 0] =  (A[:, 0] ** 2) * lambd[0]
for i in range(1, A.shape[1]):
    A_new_loop[:, i] = np.diag(
        A[:, : i + 1]
        @ np.diag(lambd[: i + 1])
        @ A[:, : i + 1].T
    ).flatten()

# 验证结果是否一致
print(np.allclose(A_new, A_new_loop))  # 输出True，说明结果完全匹配

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