如何优化基于cKDTree的多点最短路径排序代码？

现有代码的核心问题 #

1. 局部最优导致跳点：你用的是贪心最近邻算法（每次选当前点的最近未访问点），这种方法只会关注局部最短距离，必然陷入局部最优，出现绕路、跳过远端点后再回头连接的情况，也就是你看到的错误连线。
2. 性能瓶颈：find_closest_point_idx中每次查询k=len(tree.data)，相当于对所有点做全量距离排序，再遍历找第一个未访问点，数据量大时复杂度为O(n²)，速度极慢；同时用list存储已访问索引，idx not in visited_indices是O(n)查询，进一步拖慢速度。

一、解决跳点问题：改用全局路径规划算法 #

你需要的是**旅行商问题(TSP)**的近似求解（精确解对大数据量不现实），目标是找到经过所有点的最短路径，匹配你的需求。

1. 最近插入法（贪心改进版，无需额外库） #

比单纯贪心更接近全局最优，避免大量跳点：

import numpy as np
import pandas as pd
from typing import List

def nearest_insertion_path(df: pd.DataFrame, start_idx: int) -> pd.DataFrame:
    coords = df[['lat', 'lon']].values
    n = len(coords)
    visited = set([start_idx])
    path = [start_idx]

    # 初始化：找到离起点最近的点，构建初始路径
    dists = np.linalg.norm(coords - coords[start_idx], axis=1)
    dists[start_idx] = np.inf
    next_idx = np.argmin(dists)
    visited.add(next_idx)
    path.append(next_idx)

    while len(visited) < n:
        min_cost = np.inf
        insert_pos = -1
        new_idx = -1
        # 遍历所有未访问点，找插入到路径中成本最小的位置
        for idx in range(n):
            if idx in visited:
                continue
            for i in range(len(path)):
                prev_point = path[i]
                next_point = path[(i+1) % len(path)]
                # 插入成本：新增两段距离减去原路径段的距离
                cost = (np.linalg.norm(coords[prev_point]-coords[idx]) 
                        + np.linalg.norm(coords[idx]-coords[next_point]) 
                        - np.linalg.norm(coords[prev_point]-coords[next_point]))
                if cost < min_cost:
                    min_cost = cost
                    insert_pos = i + 1
                    new_idx = idx
        visited.add(new_idx)
        path.insert(insert_pos, new_idx)

    return df.iloc[path]

2. 模拟退火算法（更优的近似解，需装mlrose库） #

适合大数据量，能找到更接近全局最优的路径：

import mlrose
import numpy as np
import pandas as pd

def tsp_simulated_annealing(df: pd.DataFrame, start_idx: int) -> pd.DataFrame:
    coords = df[['lat', 'lon']].values
    n = len(coords)

    # 定义TSP适应度函数
    fitness = mlrose.TravellingSales(coords=coords)
    problem = mlrose.TSPOpt(length=n, fitness_fn=fitness, maximize=False)

    # 模拟退火求解最优路径
    best_path, _ = mlrose.simulated_annealing(problem, random_state=42)

    # 调整路径以指定起点开头
    start_pos = np.where(best_path == start_idx)[0][0]
    optimized_path = np.roll(best_path, -start_pos)

    return df.iloc[optimized_path]

二、性能优化方案 #

1. 优化最近邻查询逻辑 #

把全量k查询改成小范围查询，用布尔数组加速访问判断：

from scipy.spatial import cKDTree
from typing import Union
import numpy as np

def find_closest_point_idx(current_idx: int, tree: cKDTree, visited: np.ndarray) -> Union[int, None]:
    # 先查前10个最近点，足够找到未访问点（可根据数据调整k值）
    _, closest_indices = tree.query(tree.data[current_idx], k=10)
    for idx in closest_indices:
        if not visited[idx]:
            return idx
    # 极端情况：前10个都已访问，再查全量
    _, closest_indices = tree.query(tree.data[current_idx], k=len(tree.data))
    for idx in closest_indices:
        if not visited[idx]:
            return idx
    return None

# 主函数修改
def calculate_shortest_path_optimized(df: pd.DataFrame, starting_point: np.array) -> pd.DataFrame:
    points = df.loc[:, ['lat', 'lon']].values
    tree = cKDTree(points)
    starting_idx = tree.query(starting_point)[1]

    visited = np.zeros(len(points), dtype=bool)
    visited[starting_idx] = True
    visited_indices = [starting_idx]

    while len(visited_indices) < len(points):
        closest_idx = find_closest_point_idx(visited_indices[-1], tree, visited)
        if closest_idx is None:
            break
        visited_indices.append(closest_idx)
        visited[closest_idx] = True

    return df.iloc[visited_indices]

三、聚类精简的可行性（针对超大数据量） #

如果数据量极大，先聚类精简完全可行，针对不同类型列的处理方式：

• 数值型列：簇内取均值、中位数或加权平均（按距离权重）
• 分类型列：簇内取众数，或保留所有类别（若需后续分析）

示例代码（用DBSCAN聚类，适合地理点）：

from sklearn.cluster import DBSCAN
from sklearn.preprocessing import StandardScaler
import pandas as pd
import numpy as np

def cluster_simplify_df(df: pd.DataFrame, eps: float = 0.01, min_samples: int = 5) -> pd.DataFrame:
    # 标准化经纬度（DBSCAN对距离敏感）
    coords = StandardScaler().fit_transform(df[['lat', 'lon']])
    db = DBSCAN(eps=eps, min_samples=min_samples).fit(coords)
    df['cluster'] = db.labels_

    simplified = []
    for cluster_id in np.unique(db.labels_):
        cluster = df[df['cluster'] == cluster_id]
        if cluster_id == -1:  # 噪声点单独保留
            simplified.extend(cluster.to_dict('records'))
            continue
        # 数值型列取均值
        num_cols = cluster.select_dtypes(include=np.number).columns
        num_mean = cluster[num_cols].mean()
        # 分类型列取众数
        cat_cols = cluster.select_dtypes(exclude=np.number).columns
        cat_mode = cluster[cat_cols].mode().iloc[0]
        # 合并成代表点
        rep_point = pd.concat([num_mean, cat_mode]).to_dict()
        rep_point['cluster_size'] = len(cluster)
        simplified.append(rep_point)

    return pd.DataFrame(simplified).drop('cluster', axis=1)

总结 #

• 跳点问题本质是贪心算法的局部最优缺陷，改用TSP近似算法（最近插入法、模拟退火）可解决。
• 性能优化核心是减少全量查询，用布尔数组替代list做已访问判断。
• 超大数据量下，先聚类精简完全可行，针对不同类型列采用对应聚合方式即可。

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