如何将电厂匹配最近变电站的pandas迭代代码改写为.apply实现

注：原迭代代码存在逻辑错误：循环内直接对df1['assigned_to']赋值会覆盖所有行的结果，仅保留最后一次循环的计算值，以下方案已修复该问题。

apply实现版本 #

该版本替换了双层iterrows迭代，仅保留内层apply处理变电站距离计算，外层用apply遍历电厂数据，减少迭代开销：

import pandas as pd
import geopy.distance

# 定义单个电厂匹配最近变电站的函数
def find_nearest_ss(pp_row, ss_df):
    pp_coord = (pp_row['x'], pp_row['y'])
    # 计算当前电厂到所有变电站的测地距离
    ss_distance = ss_df.apply(
        lambda ss_row: geopy.distance.geodesic(pp_coord, (ss_row['x'], ss_row['y']).km,
        axis=1
    )
    # 获取距离最小的变电站信息
    min_idx = ss_distance.idxmin()
    return pd.Series(
        [ss_df.loc[min_idx, 'ID_ss'], ss_distance.min()],
        index=['nearest_ss_id', 'distance_km']
    )

# 对df1每行应用匹配函数，axis=1表示按行处理
df1[['nearest_ss_id', 'distance_km']] = df1.apply(
    lambda row: find_nearest_ss(row, df2),
    axis=1
)

大数量级数据优化方案 #

如果数据量超过1万条，推荐使用haversine公式向量化计算距离矩阵的方案，速度会比apply版本快10倍以上，不需要依赖geopy，示例代码如下：

import pandas as pd
import numpy as np

# 经纬度转弧度
df1[['lon_rad', 'lat_rad']] = np.radians(df1[['x', 'y']])
df2[['lon_rad', 'lat_rad']] = np.radians(df2[['x', 'y']])

R = 6371 # 地球半径，单位km

# 生成距离矩阵
lat1, lon1 = df1['lat_rad'].values, df1['lon_rad'].values
lat2, lon2 = df2['lat_rad'].values, df2['lon_rad'].values

dlat = lat2[:, np.newaxis] - lat1
dlon = lon2[:, np.newaxis] - lon1

a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2[:, np.newaxis]) * np.sin(dlon/2)**2
c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1-a))
dist_matrix = R * c

# 匹配最近变电站
min_idx = dist_matrix.argmin(axis=0)
df1['nearest_ss_id'] = df2['ID_ss'].iloc[min_idx].values
df1['distance_km'] = dist_matrix.min(axis=0)

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