基于经纬度将电厂DataFrame匹配到最近变电站的高效实现方法

一、基于apply的改造实现 #

首先修正你原有代码的几个小问题：变量a未定义、geopy的geodesic方法要求参数顺序为(纬度, 经度)，不要写反坐标导致距离计算错误。改造后的apply实现如下：

import pandas as pd
import geopy.distance

# 修正示例数据的语法错误
df1 = pd.DataFrame({'ID_pp':['p1','p2','p3','p4'],'x':[12.644881,11.563269, 12.644881,  8.153184], 'y':[48.099206, 48.020081, 48.099206, 49.153766]})
df2 = pd.DataFrame({'ID_ss':['s1','s2','s3','s4'],'x':[9.269, 9.390, 9.317, 10.061], 'y':[55.037, 54.940, 54.716, 54.349]})

# 定义匹配最近变电站的函数
def find_closest_substation(pp_row):
    # 打包当前电厂坐标，顺序为(纬度, 经度)
    pp_coord = (pp_row['y'], pp_row['x'])
    # 计算到所有变电站的测地距离
    distances = df2.apply(lambda ss_row: geopy.distance.geodesic(pp_coord, (ss_row['y'], ss_row['x'])).km, axis=1)
    # 取最小距离对应的索引
    min_idx = distances.idxmin()
    # 返回最近变电站ID和距离
    return pd.Series([df2.loc[min_idx, 'ID_ss'], distances[min_idx]], index=['closest_ss_id', 'distance_km'])

# 对每个电厂应用匹配逻辑，结果合并到df1
df1[['closest_ss_id', 'distance_km']] = df1.apply(find_closest_substation, axis=1)

该方案比双层iterrows快2~3倍，但本质还是行级迭代，数据量过万的场景仍有性能瓶颈，推荐使用下方的高性能方案。

二、更高性能的批量计算方案 #

1. 向量化Haversine公式计算 #

完全基于numpy向量运算，没有循环，速度比apply快100倍以上，适合电厂、变电站规模在几千到几万的场景：

import numpy as np

# 经纬度转弧度
df1_lat = np.radians(df1['y'])
df1_lon = np.radians(df1['x'])
df2_lat = np.radians(df2['y'])
df2_lon = np.radians(df2['x'])

# 广播计算所有电厂到所有变电站的球面距离矩阵
dlat = df2_lat.values[np.newaxis, :] - df1_lat.values[:, np.newaxis]
dlon = df2_lon.values[np.newaxis, :] - df1_lon.values[:, np.newaxis]
a = np.sin(dlat / 2)**2 + np.cos(df1_lat.values[:, np.newaxis]) * np.cos(df2_lat.values[np.newaxis, :]) * np.sin(dlon / 2)**2
c = 2 * np.arcsin(np.sqrt(a))
distance_matrix = c * 6371 # 6371为地球半径，单位为公里，结果单位为公里

# 提取每个电厂的最近变电站信息
min_indices = distance_matrix.argmin(axis=1)
min_distances = distance_matrix[np.arange(len(df1)), min_indices]
df1['closest_ss_id'] = df2['ID_ss'].iloc[min_indices].values
df1['distance_km'] = min_distances

2. KDTree空间索引方案 #

如果变电站规模超过10万，距离矩阵会占用过多内存，可使用KDTree空间索引实现低内存、毫秒级查询：

import numpy as np
from scipy.spatial import KDTree

# 经纬度（弧度）转三维球面坐标，适配KDTree计算
def lonlat_to_3d(lon, lat):
    cos_lat = np.cos(lat)
    return np.column_stack([cos_lat * np.cos(lon), cos_lat * np.sin(lon), np.sin(lat)])

# 预处理变电站坐标，构建KDTree
df2_3d = lonlat_to_3d(np.radians(df2['x'].values), np.radians(df2['y'].values))
tree = KDTree(df2_3d)

# 预处理电厂坐标，查询最近邻
df1_3d = lonlat_to_3d(np.radians(df1['x'].values), np.radians(df1['y'].values))
min_distances, min_indices = tree.query(df1_3d, k=1)
# KDTree返回的是三维弦长，转换为球面距离
min_distances = 2 * 6371 * np.arcsin(min_distances / 2)

# 合并结果
df1['closest_ss_id'] = df2['ID_ss'].iloc[min_indices].values
df1['distance_km'] = min_distances

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