如何提速双重循环构建的距离矩阵函数？（无循环优化方案）

你的矩阵元素 (M_{ij}) 本质是马氏距离的平方，数学表达式可展开为：
$$M_{ij} = (x_i - x_j)^T D^{-1} (x_i - x_j) = x_i^T D^{-1}x_i - 2x_i^T D^{-1}x_j + x_j^T D^{-1}x_j$$

基于这个展开式，我们可以用numpy的向量化运算完全替代双重循环，彻底解决Python循环的性能瓶颈。

方法1：直接矩阵求逆+广播运算 #

import numpy as np

def matrix_vectorized(X, D):
    # 计算3×3矩阵D的逆（计算成本极低）
    inv_D = np.linalg.inv(D)
    # 计算交叉项矩阵C，C[i,j] = x_i^T @ inv_D @ x_j
    C = X @ inv_D @ X.T
    # 提取对角线元素并转为列向量，用于广播计算
    diag_C = np.diag(C)[:, np.newaxis]
    # 用广播生成完整M矩阵，最后扁平化输出（与原函数返回格式一致）
    M = diag_C + diag_C.T - 2 * C
    return M.flatten()

方法2：Cholesky分解（数值稳定性更优） #

如果D是正定矩阵（距离函数通常要求D正定），推荐用Cholesky分解替代直接求逆，数值稳定性更好：

def matrix_vectorized_cholesky(X, D):
    # 对D做Cholesky分解，得到下三角矩阵L，满足D = L @ L.T
    L = np.linalg.cholesky(D)
    # 将X转换到新空间：Y[i] = x_i @ L^{-1}
    Y = X @ np.linalg.inv(L)
    # 交叉项矩阵C = Y @ Y.T，等价于x_i^T @ D^{-1} @ x_j
    C = Y @ Y.T
    diag_C = np.diag(C)[:, np.newaxis]
    M = diag_C + diag_C.T - 2 * C
    return M.flatten()

性能对比 #

针对你给出的1000×3的X和3×3的D：

• 原双重循环函数：需要执行1,000,000次np.linalg.solve调用，Python循环本身会带来极大开销。
• 向量化版本：仅需几次矩阵乘法运算，底层由优化过的BLAS/LAPACK库实现，速度能提升100倍以上。

验证正确性 #

用小规模数据验证两种方法的结果一致性：

# 生成测试数据
X_test = np.random.rand(5, 3)
D_test = np.random.rand(3, 3)
D_test = D_test @ D_test.T  # 构造正定矩阵

# 计算结果
M_loop = matrix(X_test, D_test)
M_vec = matrix_vectorized(X_test, D_test)

# 验证误差在浮点精度范围内
print(np.allclose(M_loop, M_vec))  # 输出True

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