如何高效计算两个DataFrame经纬度距离并筛选符合阈值的a表数据

嵌套for循环效率低的核心原因是Python原生循环开销大，且时间复杂度为O(n*m)（n为a的点数、m为b的点数），数据量稍大就会严重卡顿，以下是两种高效实现方案：

方案1：KD树近邻搜索（推荐，大数据量最优） #

基于空间索引实现，时间复杂度可降至O(n log m)，支持多线程加速，适合千级以上点数的场景，使用scikit-learn的NearestNeighbors实现：

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

# 距离阈值，单位米，根据实际需求修改
X = 500

# 1. 经纬度转弧度（haversine度量要求输入为弧度）
a['lat_rad'] = np.deg2rad(a['LAT'])
a['lon_rad'] = np.deg2rad(a['LONG'])
b['lat_rad'] = np.deg2rad(b['LAT'])
b['lon_rad'] = np.deg2rad(b['LONG'])

# 2. 用b的点位构建半径搜索KD树
R = 6371000 # 地球平均半径，单位米
threshold_rad = X / R
nbrs = NearestNeighbors(radius=threshold_rad, metric='haversine', n_jobs=-1)
nbrs.fit(b[['lat_rad', 'lon_rad']].values)

# 3. 对a的所有点位做半径查询，保留存在至少一个近邻的点位
_, indices = nbrs.radius_neighbors(a[['lat_rad', 'lon_rad']].values)
mask = [len(idx) > 0 for idx in indices]
a_filtered = a[mask].drop(columns=['lat_rad', 'lon_rad']).copy()

方案2：向量化距离计算（小数据量最简） #

基于numpy广播机制计算所有点对距离，不需要额外引入机器学习依赖，适合百级点数的场景：

import pandas as pd
import numpy as np

X = 500 # 距离阈值，单位米

# 1. 经纬度转弧度
a_lat, a_lon = np.deg2rad(a['LAT']), np.deg2rad(a['LONG'])
b_lat, b_lon = np.deg2rad(b['LAT']), np.deg2rad(b['LONG'])

# 2. 广播计算所有点对的haversine距离（单位米）
dlat = b_lat[np.newaxis, :] - a_lat[:, np.newaxis]
dlon = b_lon[np.newaxis, :] - a_lon[:, np.newaxis]
a_ = np.sin(dlat / 2)**2 + np.cos(a_lat[:, np.newaxis]) * np.cos(b_lat[np.newaxis, :]) * np.sin(dlon / 2)**2
c = 2 * np.arcsin(np.sqrt(a_))
dist_matrix = 6371000 * c

# 3. 筛选出任意距离小于阈值的a点位
mask = (dist_matrix < X).any(axis=1)
a_filtered = a[mask].copy()

高精度需求补充方案 #

如果需要和geopy.distance.distance的WGS84椭球计算结果完全对齐，可使用pyproj的向量化计算：

import pandas as pd
import numpy as np
from pyproj import Geod

X = 500 # 距离阈值，单位米

# 初始化WGS84椭球计算器
g = Geod(ellps='WGS84')

# 广播展开所有点对坐标
lons1 = a['LONG'].values[:, np.newaxis].repeat(len(b), axis=1)
lats1 = a['LAT'].values[:, np.newaxis].repeat(len(b), axis=1)
lons2 = b['LONG'].values[np.newaxis, :].repeat(len(a), axis=0)
lats2 = b['LAT'].values[np.newaxis, :].repeat(len(a), axis=0)

# 批量计算距离，返回结果单位为米
_, _, dist_matrix = g.inv(lons1, lats1, lons2, lats2)

mask = (dist_matrix < X).any(axis=1)
a_filtered = a[mask].copy()

以上所有方案的底层运算均为C实现，比原生Python嵌套循环效率高10~1000倍不等，具体提升幅度随数据量增大而升高。

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