Pandas按Group列分组计算连续行的haversine地理距离

原有实现 #

使用如下公式计算得到对应DataFrame的dist列：

df['dist'] = haversine(df.LAT.shift(), df.LONG.shift(),df.loc[1:, 'LAT'], df.loc[1:, 'LONG'])

原有输出结果 #

Group       ID      LAT       LONG         dist
   1         1  74.166061  30.512811          NaN
   1         2  72.249672  33.427724   232.549785
   1         3  67.499828  37.937264   554.905446
   1         4  84.253715  69.328767  1981.896491
   2         5  72.104828  33.823462  1513.397997
   2         6  63.989462  51.918173  1164.481327
   2         7  80.209112  33.530778  1887.256899
   2         8  68.954132  35.981256  1252.531365
   2         9  83.378214  40.619652  1606.340727
   2        10  68.778571   6.607066  1793.921854

需求说明 #

改写计算逻辑，按Group列分组计算，每个分组内仅计算相邻行的距离，分组首行的dist值为NaN，期望输出如下：

Group       ID      LAT       LONG         dist
   1         1  74.166061  30.512811          NaN
   1         2  72.249672  33.427724   232.549785
   1         3  67.499828  37.937264   554.905446
   1         4  84.253715  69.328767  1981.896491
   2         5  72.104828  33.823462          NaN
   2         6  63.989462  51.918173  1164.481327
   2         7  80.209112  33.530778  1887.256899
   2         8  68.954132  35.981256  1252.531365
   2         9  83.378214  40.619652  1606.340727
   2        10  68.778571   6.607066  1793.921854

实现代码 #

首先给出haversine函数的完整实现，再给出分组计算的核心逻辑：

import pandas as pd
import numpy as np

# 经纬度距离计算函数：默认返回距离单位为千米
def haversine(lat1, lon1, lat2, lon2, unit='km'):
    # 角度转弧度
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    # 计算差值
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    # 代入公式计算
    a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    # 地球半径：千米为单位取6371，英里为单位取3956
    r = 6371 if unit == 'km' else 3956
    return c * r

# 分组计算dist列核心代码
df['dist'] = df.groupby('Group').apply(
    lambda group: haversine(group['LAT'].shift(), group['LONG'].shift(), group['LAT'], group['LONG'])
).reset_index(level=0, drop=True)

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