如何快速计算numpy矩阵的高斯核距离矩阵以替代嵌套循环

我们可以利用L2范数平方的代数展开式完全避免嵌套循环，用底层C实现的向量化运算完成计算，速度比纯Python嵌套循环快3~4个数量级。

纯numpy向量化实现（适用于中小规模数据） #

核心原理：L2范数平方可以展开为 ||X[i] - X[j]||² = ||X[i]||² + ||X[j]||² - 2<X[i], X[j]>，结合numpy的广播机制可以一次性计算所有位置的距离平方值。

import numpy as np

X = np.array([[1,2,3],[4,5,6],[7,8,9]])
# 计算每行的L2范数平方
row_norms = np.sum(X ** 2, axis=1)
# 广播计算全量成对距离平方
dist_sq = row_norms[:, np.newaxis] + row_norms - 2 * X @ X.T
# 生成目标矩阵P
P = np.exp(-dist_sq)

print(P)

运行输出和朴素嵌套循环实现完全一致。

scipy优化实现（速度更快、代码更简洁） #

scipy内置的成对距离计算函数pdist经过专门的性能优化，比手写numpy广播实现速度更高，内存占用也更低：

import numpy as np
from scipy.spatial.distance import pdist, squareform

X = np.array([[1,2,3],[4,5,6],[7,8,9]])
# 直接计算所有成对欧氏距离的平方，转为方阵格式
dist_sq = squareform(pdist(X, metric='sqeuclidean'))
P = np.exp(-dist_sq)

print(P)

5e4量级大规模数据注意事项 #

• 直接生成全量[5e4, 5e4]矩阵的话，共25亿个元素，用float64存储需要20GB内存，float32也需要10GB。内存不足的场景可以采用分块计算：每次只计算k行对应的k*m大小的子矩阵，计算完成后直接写入磁盘或做后续处理，避免全量加载到内存。
• 高斯核数值随距离增大快速衰减，距离超过一定阈值的元素数值极小，业务允许的前提下可以直接截断为0，用稀疏矩阵格式存储P，可降低数倍到数十倍的内存占用。

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