Python优化30万行与3k行DataFrame的Haversine距离匹配效率

针对30万行数据逐行遍历找最近点的性能瓶颈，核心优化思路是用空间索引（如BallTree/KDTree）将O(n*m)的时间复杂度降到O(n log m)，以下是具体实现方案：

方案一：用sklearn BallTree直接处理球面距离 #

BallTree天生适合处理高维空间的最近邻查询，且支持haversine度量直接计算球面距离，无需手动转换坐标：

import numpy as np
from sklearn.neighbors import BallTree
import pandas as pd

# 1. 预处理坐标：将经纬度转为弧度（haversine要求输入为弧度）
# df是3k行的网格子集，提取[纬度, 经度]并转弧度
grid_coords = np.radians(df[['Lat', 'Lon']].values)
# 构建BallTree，指定度量为haversine（球面距离）
tree = BallTree(grid_coords, metric='haversine')

# 2. 处理30万行的暴露数据：同样转弧度
exposure_coords = np.radians(df_exposure[['lat_loc', 'lon_loc']].values)
# 查询每个暴露点的最近网格点，返回距离（弧度）和对应网格索引
_, nearest_grid_indices = tree.query(exposure_coords, k=1)
# 展平索引数组（query返回的是二维数组）
nearest_grid_indices = nearest_grid_indices.flatten()

# 3. 累加Limit到对应网格点
# 将暴露数据的Limit按最近网格索引分组求和
total_exposure = df_exposure.groupby(nearest_grid_indices)['Limit'].sum()
# 将结果合并回原网格DataFrame
df['total_limit'] = total_exposure.values

方案二：用GeoPandas空间索引（适合地理数据场景） #

如果本身是地理数据处理流程，GeoPandas的空间索引（R-tree）更直观，适合后续扩展地理操作：

import geopandas as gpd
from shapely.geometry import Point

# 1. 转换为GeoDataFrame，指定WGS84坐标系（EPSG:4326）
gdf_exposure = gpd.GeoDataFrame(
    df_exposure,
    geometry=gpd.points_from_xy(df_exposure.lon_loc, df_exposure.lat_loc),
    crs="EPSG:4326"
)

gdf_grid = gpd.GeoDataFrame(
    df,
    geometry=gpd.points_from_xy(df.Lon, df.Lat),
    crs="EPSG:4326"
)

# 2. 构建网格数据的空间索引
grid_sindex = gdf_grid.sindex

# 3. 批量查询最近点（空间索引先缩小候选范围，再计算精确距离）
def get_nearest_grid_idx(point):
    # 用空间索引快速获取候选点（避免全量计算）
    candidate_indices = list(grid_sindex.nearest(point.bounds, 1))[0]
    # 在候选点中计算精确距离，找到最近的索引
    candidate_distances = gdf_grid.iloc[candidate_indices].distance(point)
    return candidate_distances.idxmin()

# 批量处理所有暴露点
gdf_exposure['nearest_grid_idx'] = gdf_exposure.geometry.apply(get_nearest_grid_idx)

# 4. 累加Limit
total_exposure = gdf_exposure.groupby('nearest_grid_idx')['Limit'].sum()
df['total_limit'] = total_exposure.values

为什么循环慢？ #

Python的逐行循环属于解释型操作，30万行×3k次距离计算=9亿次操作，开销极大。而空间索引基于底层C/C++实现的向量运算，将查询复杂度从O(n*m)降至O(n log m)，性能提升可达几十甚至上百倍。

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