如何用Numpy/Scipy高效实现矩阵列向量间的欧氏距离矩阵？

一、Scipy内置函数直接实现 #

Scipy的scipy.spatial.distance.cdist专门用于计算两个点集之间所有点对的距离，默认就是欧氏距离，完全匹配需求。需要注意的是，cdist要求输入矩阵形状为**(样本数, 特征维度)**，而你的矩阵A是(d,n)、B是(d,m)，所以需要先转置为(n,d)和(m,d)。

代码示例：

import numpy as np
from scipy.spatial.distance import cdist

def distance_matrix_scipy(A, B):
    return cdist(A.T, B.T)  # 默认计算欧氏距离

验证示例：

A = np.array([[-1, -1, 1, 1, 2], [1, -1, 2, -1, 1]])
B = np.array([[-2, -1, 1, 2], [-1, 2, 1, -1]])

print(distance_matrix_scipy(A, B))

输出和原函数的结果完全一致。

二、Numpy向量化实现 #

利用欧氏距离的数学公式展开实现无循环的向量化计算，比双重循环效率提升显著，尤其适合大规模数据：

欧氏距离平方公式：$|a - b|^2 = |a|^2 + |b|^2 - 2a \cdot b$

对公式开根号即可得到欧氏距离，结合Numpy的广播机制可以快速计算：

代码示例：

import numpy as np

def distance_matrix_vectorized(A, B):
    # 计算A各列的范数平方 (n,)
    norm_A = np.sum(A**2, axis=0)
    # 计算B各列的范数平方 (m,)
    norm_B = np.sum(B**2, axis=0)
    # 利用广播计算所有点对的距离平方矩阵，再开根号
    dist_sq = norm_A[:, np.newaxis] + norm_B[np.newaxis, :] - 2 * A.T @ B
    # 防止数值误差导致的极小负数，取最大值后开根号
    return np.sqrt(np.maximum(dist_sq, 0))

验证示例：

print(distance_matrix_vectorized(A, B))

输出同样与原函数结果一致。

效率对比 #

• 原双重循环的时间复杂度为O(nmd)，且循环在Python层面执行，速度最慢；
• Numpy向量化实现的时间复杂度为O(nm + nd + m*d)，所有运算在C层面执行，效率大幅提升；
• Scipy的cdist底层也是优化后的C实现，效率与Numpy向量化相当，且支持多种距离度量，使用更便捷。

内容的提问来源于stack exchange，提问作者jan-seins