如何加速Pandas中两个DataFrame逐行求最短距离的计算

你的问题核心是避免O(n*m)的双重循环，改用O(n log m)的高效最近邻搜索方法，以下是几种实用的加速方案：

方案1：Scipy KDTree（最快最推荐） #

利用Scipy的KDTree实现球面距离的快速最近邻查询，适合大规模数据匹配：

import pandas as pd
from scipy.spatial import KDTree
import numpy as np

# 加载你的DataFrame（示例假设已完成加载）
# table1 = pd.read_csv("table1.csv")
# table2 = pd.read_csv("table2.csv")

# 将经纬度转换为弧度（Haversine距离要求输入弧度）
table2_rad = np.radians(table2[["lat", "long"]].values)
table1_rad = np.radians(table1[["lat", "long"]].values)

# 构建KDTree，使用Haversine度量（对应球面距离）
kdtree = KDTree(table2_rad, metric="haversine")

# 查询每个table1点的最近邻，返回距离（弧度）和对应table2的索引
_, nearest_indices = kdtree.query(table1_rad, k=1)

# 填充table1的table2_id列
table1["table2_id"] = table2.iloc[nearest_indices]["table2_id"].values

优势 #

• 时间复杂度为O(n log m)，4万行table1匹配3千行table2，几秒内即可完成
• 内存占用低，无需生成庞大的距离矩阵

方案2：Geopandas空间索引（地理数据专业方案） #

如果需要更贴合地理数据的处理逻辑，用Geopandas的空间索引先做粗过滤，再精确计算距离：

import geopandas as gpd
from shapely.geometry import Point

# 转换为GeoDataFrame，创建点几何对象
table1_gdf = gpd.GeoDataFrame(
    table1,
    geometry=gpd.points_from_xy(table1["long"], table1["lat"])
)
table2_gdf = gpd.GeoDataFrame(
    table2,
    geometry=gpd.points_from_xy(table2["long"], table2["lat"])
)

# 为table2构建空间索引
sindex = table2_gdf.sindex

# 定义函数获取最近的table2_id
def get_nearest_id(point):
    # 空间索引快速筛选候选点（减少计算量）
    candidate_idx = list(sindex.nearest((point.x, point.y), 1))
    candidates = table2_gdf.iloc[candidate_idx]
    # 计算精确距离并找到最近点
    nearest_row = candidates.distance(point).idxmin()
    return table2_gdf.loc[nearest_row, "table2_id"]

# 批量填充
table1["table2_id"] = table1_gdf.geometry.apply(get_nearest_id)

优势 #

• 空间索引先缩小候选范围，避免全量距离计算
• 支持更复杂的地理操作，后续扩展更方便

方案3：分块计算（临时过渡方案） #

如果暂时无法使用上述库，可通过分块减少内存压力，比纯双重循环快：

from geopy.distance import geodesic
import numpy as np

chunk_size = 1000  # 每次处理1000行，可根据内存调整
table1["table2_id"] = np.nan

for start in range(0, len(table1), chunk_size):
    end = min(start + chunk_size, len(table1))
    chunk = table1.iloc[start:end]
    
    # 广播计算当前块与table2所有点的距离
    distances = np.array([
        [geodesic((lat1, lon1), (lat2, lon2)).meters 
         for lat2, lon2 in zip(table2["lat"], table2["long"])]
        for lat1, lon1 in zip(chunk["lat"], chunk["long"])
    ])
    
    # 找到每个点的最近邻索引并填充
    min_indices = distances.argmin(axis=1)
    table1.loc[start:end-1, "table2_id"] = table2.iloc[min_indices]["table2_id"].values

说明 #

• 仍为O(n*m)复杂度，但分块减少了单次内存占用
• 速度比纯循环快，但远不如前两种方案

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