如何用向量化优化Numpy数组行间的欧氏距离计算？

原双层Python循环的效率瓶颈在于Python循环本身的开销，以及频繁调用np.linalg.norm的函数调用成本。利用Numpy的广播机制和矩阵运算，我们可以把循环逻辑转化为底层C实现的批量运算，大幅提升运行速度。以下是几种高效的向量化实现方案：

方案1：基于欧氏距离公式展开（最优性能+低内存） #

欧氏距离的平方可以展开为：$|a-b|^2 = |a|^2 + |b|^2 - 2a \cdot b$，我们可以用矩阵运算直接计算整个距离矩阵，最后开平方得到结果：

import numpy as np

# 计算A每行的L2范数平方，保持维度便于广播
norm_A = np.sum(A ** 2, axis=1, keepdims=True)  # 形状 (n1, 1)
# 计算B每行的L2范数平方并转置，匹配广播维度
norm_B = np.sum(B ** 2, axis=1, keepdims=True).T  # 形状 (1, n2)
# 计算A与B的点积矩阵
dot_product = A @ B.T  # 形状 (n1, n2)
# 计算欧氏距离矩阵，用maximum避免数值精度问题导致的负数
C = np.sqrt(np.maximum(norm_A + norm_B - 2 * dot_product, 0))

说明： #

• 这种方法完全依赖矩阵运算，没有中间大数组生成，内存占用最低，速度最快，尤其适合大样本量的场景。
• 加入np.maximum是为了避免因浮点精度误差出现极小负数，导致开平方报错。

方案2：基于广播直接计算差值范数（直观易理解） #

通过扩展数组维度，让Numpy自动广播实现两两行的差值计算，再批量计算范数：

import numpy as np

# 扩展A的维度为(n1, 1, m)，B的维度为(1, n2, m)
A_expanded = A[:, np.newaxis, :]
B_expanded = B[np.newaxis, :, :]
# 对第三维度（特征维度）计算L2范数，得到(n1, n2)的距离矩阵
C = np.linalg.norm(A_expanded - B_expanded, axis=2)

说明： #

• 逻辑和原循环完全对应，非常直观，但会生成一个形状为(n1, n2, m)的中间差值数组，当n1、n2、m都很大时，内存开销会显著增加。

方案3：使用Scipy现成函数（最简洁） #

如果项目允许依赖Scipy，scipy.spatial.distance.cdist是专门用于计算两两距离的优化函数，代码最简洁：

from scipy.spatial.distance import cdist

# 直接计算A和B每行的欧氏距离
C = cdist(A, B, metric='euclidean')

说明： #

• 该函数内部做了高度优化，支持多种距离度量（如曼哈顿距离、余弦距离等），无需自己实现公式，适合快速开发场景。

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