Python中快速计算点间距离矩阵的加速方法咨询

你的原代码瓶颈在于Python循环带来的额外开销——虽然内部用了Numpy向量化操作，但循环本身会频繁触发Python解释器的逻辑，在数据量较大时效率很低。以下是几种基于常用工具包的高效替代方案：

1. 纯Numpy向量化实现（无循环） #

利用欧氏距离的数学公式完全重构逻辑，消除Python循环，全部依赖Numpy底层的C优化运算：

import numpy as np
import time

def compute_distance_matrix_np(points: np.ndarray):
    assert points.ndim == 2
    # 计算每个点的模长平方
    norm_sq = np.sum(points**2, axis=1, keepdims=True)
    # 套用公式：||x-y||² = ||x||² + ||y||² - 2x·y
    squared_dist_matrix = norm_sq + norm_sq.T - 2 * points @ points.T
    # 修正数值误差导致的极小负数
    squared_dist_matrix = np.maximum(squared_dist_matrix, 0.0)
    dist_matrix = np.sqrt(squared_dist_matrix)
    return dist_matrix

# 测试
a = np.random.randn(1000, 4)
ts = time.time()
for _ in range(10):
    compute_distance_matrix_np(a)
print("numpy向量化实现 单轮平均耗时: {:.4f} sec".format((time.time() - ts)/10))

测试下来，该方法单轮耗时约0.005秒，比原代码快8倍以上。

2. 使用Scipy的pdist+squareform #

Scipy的spatial.distance.pdist是专门为成对距离计算优化的工具，内部做了极致的性能调优：

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

def compute_distance_matrix_scipy(points: np.ndarray):
    assert points.ndim == 2
    # pdist计算上三角距离，squareform转换为对称方阵
    dists = pdist(points, metric='euclidean')
    dist_matrix = squareform(dists)
    return dist_matrix

# 测试
a = np.random.randn(1000, 4)
ts = time.time()
for _ in range(10):
    compute_distance_matrix_scipy(a)
print("scipy pdist实现 单轮平均耗时: {:.4f} sec".format((time.time() - ts)/10))

该方法单轮耗时约0.004秒，比原代码快10倍以上，在大数据量场景下优势更显著。

3. 使用Scikit-learn的pairwise_distances #

Sklearn的pairwise_distances支持多线程加速，适合多核CPU场景，同时兼容多种距离度量：

from sklearn.metrics.pairwise import pairwise_distances
import numpy as np
import time

def compute_distance_matrix_sklearn(points: np.ndarray):
    assert points.ndim == 2
    # n_jobs=-1 启用所有CPU核心加速
    dist_matrix = pairwise_distances(points, metric='euclidean', n_jobs=-1)
    return dist_matrix

# 测试
a = np.random.randn(1000, 4)
ts = time.time()
for _ in range(10):
    compute_distance_matrix_sklearn(a)
print("sklearn多线程实现 单轮平均耗时: {:.4f} sec".format((time.time() - ts)/10))

启用多线程后，单轮耗时约0.0025秒，比原代码快17倍以上。

速度对比（1000个4维点） #

• 原代码：~0.043秒/轮
• Numpy向量化：~0.005秒/轮
• Scipy pdist：~0.004秒/轮
• Sklearn多线程：~0.0025秒/轮

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