如何加速Python中点的最近线性距离匹配及Color映射代码？

问题分析 #

原代码采用双重嵌套循环，时间复杂度为O(N*M)，其中N=572，M=1e6，总计算量超5亿次，再加上循环内的print IO操作，导致运行速度极慢。另外，手动计算距离时的平方根操作属于冗余计算——比较距离大小只需平方距离即可，无需开根号。

优化方案 #

使用KDTree最近邻搜索（基于scipy库），将时间复杂度降至O(M logM + N logM)，同时移除冗余计算和低效IO操作。

步骤1：导入依赖库 #

import numpy as np
from scipy.spatial import KDTree

步骤2：构建KDTree并执行最近邻搜索 #

# 确保输入为numpy数组（原代码中已满足索引操作，可省略此步）
X = np.asarray(X)
Velocity_Momentum = np.asarray(Velocity_Momentum)
Color = np.asarray(Color)

# 构建KDTree，用于高效最近邻查询
tree = KDTree(Velocity_Momentum)

# 查询X中每个点的最近邻索引（k=1表示只取最近的1个点）
_, nearest_indices = tree.query(X, k=1)

# 直接通过索引获取对应Color值
Outputs = Color[nearest_indices]

步骤3：批量输出结果（替代循环内print） #

循环内的print会频繁触发IO操作拖慢速度，建议改为批量输出：

for x2, y2, color in zip(X[:, 0], X[:, 1], Outputs):
    print(f"X2 = {x2}  Y2 = {y2}  Color = {color}")

优化效果说明 #

• 时间复杂度大幅降低：KDTree构建时间为O(M logM)，单条查询时间为O(logM)，总计算量从5.72e8次降至约2000万次，速度提升数十倍以上。
• 移除冗余计算：KDTree内部采用优化的距离比较逻辑，无需手动计算平方根。
• 减少IO开销：批量输出避免了循环内频繁IO的性能损耗。

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