加权欧氏距离的高效矩阵构建方法问询

我在二维欧氏空间中有M个点，存储在M×2的数组X中。需要构建代价矩阵，其中元素(i,j)为距离d(X[i, :], X[j, :])，距离函数是矩阵D的逆加权的欧氏距离，公式为：

d(x,y) = <D⁻¹(x-y), x-y>

我已经尽量避免使用for循环，但想知道有没有更高效的实现方式，当前代码如下：

import numpy as np

Dinv = np.linalg.inv(D)

def cost(X, Dinv):
    Msq = len(X) ** 2
    mesh = []
    for i in range(2):  # 拆分每个坐标轴
        xmesh = np.meshgrid(X[:, i], X[:, i])  # 生成坐标轴网格
        xmesh = xmesh[1] - xmesh[0]  # 计算差值矩阵
        xmesh = xmesh.reshape(Msq)  # 重塑为向量
        mesh.append(xmesh)  # 添加到列表中

    meshv = np.vstack((mesh[0], mesh[1])).T  # 重新组合坐标差

    # 应用D^{-1}
    Dx = np.einsum("ij,kj->ki", Dinv, meshv)

    return np.sum(Dx * meshv, axis=1)  # 计算点积

更高效的实现方案 #

我们可以通过数学公式展开彻底规避meshgrid和循环，直接利用numpy的矩阵运算批量计算，在效率和内存占用上都会更优：

公式推导

把距离公式展开：
[
d(x,y) = (x-y)^T D^{-1} (x-y) = x^T D^{-1}x + y^T D^{-1}y - 2x^T D^{-1}y
]
这个展开式可以直接用矩阵乘法和广播来批量计算所有两两组合的距离。

优化后的代码

import numpy as np

def optimized_cost(X, Dinv):
    # 计算每个点的 x^T D^{-1} x
    xDinvx = np.einsum('ij,ji->i', X, Dinv @ X.T)
    # 利用广播生成完整代价矩阵，再扁平化（和原函数输出格式一致）
    cost_matrix = xDinvx[:, None] + xDinvx[None, :] - 2 * (X @ Dinv @ X.T)
    return cost_matrix.ravel()

优势说明

• 内存更高效：原方法需要生成M²×2的meshv数组，优化方法仅需处理M×M级别的矩阵，当M很大时内存差距明显。
• 计算更快：矩阵乘法和einsum都是numpy底层优化的操作，比手动拆分坐标轴的方式运算效率更高。
• 代码更简洁：直接对应数学公式，逻辑清晰，可读性更强。

验证一致性

可以用小数据集测试两种方法的结果是否一致：

# 测试示例
D = np.array([[2, 0], [0, 3]])
Dinv = np.linalg.inv(D)
X = np.array([[1,2], [3,4], [5,6]])

original_result = cost(X, Dinv)
optimized_result = optimized_cost(X, Dinv)

print(np.allclose(original_result, optimized_result))  # 输出True，说明结果一致

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