如何无需计算完整numpy矩阵乘积即可高效提取其对角元素

核心原理 #

当前实现需要先计算完整的m×n规模的乘积矩阵再提取对角线，大部分计算和内存开销都浪费在了非对角线元素上。我们可以通过代数等价转换，直接计算目标对角元素，完全跳过无关值的计算和存储。

优化实现 #

方案1：使用einsum实现（优先推荐） #

numpy的einsum支持按索引规则直接计算目标元素，开销极低：

import numpy as np

# 示例参数
m = int(1e6)
n = int(1e2)
A = np.random.random((m, n))
B = np.random.random((n, n))

# 预计算B@B.T，仅n×n规模，开销可忽略
bbT = B @ B.T
# 直接计算对角元素，不需要生成完整乘积矩阵
opt_result = np.einsum('ij,ji->i', A, bbT.T)[:min(m, n)]

方案2：显式循环实现（更易理解） #

如果对einsum的语法不熟悉，也可以通过循环直接取对应行和列点积得到结果，性能略低于einsum但远好于原始实现：

bbT = B @ B.T
opt_result = np.array([A[i] @ bbT[:, i] for i in range(min(m, n))])

正确性校验 #

可以用小规模矩阵验证优化结果和原始结果一致：

# 小规模测试
m_test = 200
n_test = 10
A_test = np.random.random((m_test, n_test))
B_test = np.random.random((n_test, n_test))

orig_result = np.diag(A_test @ B_test @ B_test.T)
print(np.allclose(orig_result, opt_result)) # 输出为True表示结果一致

开销对比 #

在m远大于n的场景下，优化后的内存开销可以降低几个数量级。

补充说明 #

如果你实际需要的是A的每行对应的二次型结果（即结果长度为m，对应表达式为np.diag(A @ B @ B.T @ A.T)，原始写法存在笔误），可以用以下einsum实现，避免生成无法存储的m×m规模矩阵：

bbT = B @ B.T
opt_result = np.einsum('ij,jk,ik->i', A, bbT, A)

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