为百万级df_base添加df_lookup中最近经纬度列的性能优化需求

原代码慢的核心原因是逐行嵌套循环：对df_base的100万行数据，每行都要遍历df_lookup的1万行计算距离，总计算量达到1e10次，时间复杂度为O(M*N)，完全无法支撑大规模数据。以下是几个高效的优化方案，按实用性排序：

方案1：Scipy KDTree（推荐，内存友好+速度极快） #

KDTree是空间索引结构，能将最近邻查询的时间复杂度降到O(M logN)，完美适配你的数据规模。

import numpy as np
from scipy.spatial import KDTree
from math import radians

# 统一lookup的列名，避免后续混乱
df_lookup = df_lookup.rename(columns={'lat.': 'lat', 'long': 'lon'})

# 1. 将经纬度转换为弧度（符合haversine公式要求）
lookup_rad = np.array([df_lookup['lat'].apply(radians), df_lookup['lon'].apply(radians)]).T
base_rad = np.array([df_base['Latitude'].apply(radians), df_base['Longitude'].apply(radians)]).T

# 2. 构建KDTree空间索引
kdtree = KDTree(lookup_rad)

# 3. 批量查询每个base点的最近邻（返回距离和对应lookup的索引）
_, nearest_idx = kdtree.query(base_rad, k=1)

# 4. 匹配对应的经纬度到df_base
df_base['lookup lat'] = df_lookup.iloc[nearest_idx]['lat'].values
df_base['lookup long'] = df_lookup.iloc[nearest_idx]['lon'].values

# 处理空值：对Latitude/Longitude为空的行设为NaN
valid_mask = df_base[['Latitude', 'Longitude']].notnull().all(axis=1)
df_base.loc[~valid_mask, ['lookup lat', 'lookup long']] = np.nan

方案2：Geopandas sjoin_nearest（代码简洁，地理数据专用） #

Geopandas基于R-tree空间索引实现了原生的最近邻空间连接，代码更直观，无需手动处理弧度转换。

import geopandas as gpd
from shapely.geometry import Point

# 1. 将普通DataFrame转为GeoDataFrame（指定WGS84坐标系）
gdf_base = gpd.GeoDataFrame(
    df_base,
    geometry=gpd.points_from_xy(df_base['Longitude'], df_base['Latitude']),
    crs="EPSG:4326"
)

# 处理lookup数据，统一列名并转为GeoDataFrame
df_lookup = df_lookup.rename(columns={'lat.': 'lat', 'long': 'lon'})
gdf_lookup = gpd.GeoDataFrame(
    df_lookup,
    geometry=gpd.points_from_xy(df_lookup['lon'], df_lookup['lat']),
    crs="EPSG:4326"
)

# 2. 执行最近邻空间连接，保留base所有行
gdf_joined = gpd.sjoin_nearest(gdf_base, gdf_lookup, how='left')

# 3. 提取结果列
df_base[['lookup lat', 'lookup long']] = gdf_joined[['lat', 'lon']]

# 处理空值
valid_mask = df_base[['Latitude', 'Longitude']].notnull().all(axis=1)
df_base.loc[~valid_mask, ['lookup lat', 'lookup long']] = np.nan

方案3：向量化Haversine计算（内存压力较大，适合lookup行数较少的场景） #

用numpy广播实现全量向量化计算，避免Python循环，比原apply快几十倍，但1万*100万的数组会占用约8GB内存（float64类型），内存充足时可尝试。

import numpy as np
from math import radians

# 转为numpy数组提升计算效率
base_lat = df_base['Latitude'].values
base_lon = df_base['Longitude'].values
lookup_lat = df_lookup['lat.'].values
lookup_lon = df_lookup['long'].values

# 转换为弧度
base_lat_rad = np.radians(base_lat)
base_lon_rad = np.radians(base_lon)
lookup_lat_rad = np.radians(lookup_lat)
lookup_lon_rad = np.radians(lookup_lon)

# 广播计算经纬度差值
dlat = lookup_lat_rad[np.newaxis, :] - base_lat_rad[:, np.newaxis]
dlon = lookup_lon_rad[np.newaxis, :] - base_lon_rad[:, np.newaxis]

# 向量化计算haversine距离
a = np.sin(dlat/2)**2 + np.cos(base_lat_rad[:, np.newaxis]) * np.cos(lookup_lat_rad[np.newaxis, :]) * np.sin(dlon/2)**2
c = 2 * np.arcsin(np.sqrt(a))
distances = 6371 * c

# 找到每个base点的最近lookup索引
nearest_idx = np.nanargmin(distances, axis=1)

# 匹配结果
df_base['lookup lat'] = lookup_lat[nearest_idx]
df_base['lookup long'] = lookup_lon[nearest_idx]

# 处理空值
valid_mask = df_base[['Latitude', 'Longitude']].notnull().all(axis=1)
df_base.loc[~valid_mask, ['lookup lat', 'lookup long']] = np.nan

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