Pandas中按组计算滚动差值并除以往期组总和的优化实现

我希望找到更简便高效的实现方法（最好使用链式调用提升代码可读性），计算按城市分组的Pop滚动差值除以往期对应日期的总Pop，最终得到结果表中的pc列。以下是现有实现及优化方案：

原始数据 #

import pandas as pd

df = pd.DataFrame(
    {
        "Date": ["2020-01-01", "2020-01-01", "2020-01-01", "2021-01-01", "2021-01-01", "2021-01-01", "2022-01-01", "2022-01-01", "2022-01-01"],
        "City": ["London", "New York", "Tokyo", "London", "New York", "Tokyo", "London", "New York", "Tokyo"],
        "Pop": [90, 70, 60, 85, 60, 45, 70, 40, 32],
    }
)

数据预览：

Date      City  Pop
0  2020-01-01    London   90
1  2020-01-01  New York   70
2  2020-01-01     Tokyo   60
3  2021-01-01    London   85
4  2021-01-01  New York   60
5  2021-01-01     Tokyo   45
6  2022-01-01    London   70
7  2022-01-01  New York   40
8  2022-01-01     Tokyo   32

现有实现代码 #

df['pop_diff'] = df.groupby(['City'])['Pop'].diff()
df['total'] = df.groupby('Date').Pop.transform('sum')
df['total_shift'] = df.groupby('City')['total'].shift()
df['pc'] = df['pop_diff'] / df['total_shift']

计算结果：

Date      City  Pop  pop_diff  total  total_shift        pc
0  2020-01-01    London   90       NaN    220          NaN       NaN
1  2020-01-01  New York   70       NaN    220          NaN       NaN
2  2020-01-01     Tokyo   60       NaN    220          NaN       NaN
3  2021-01-01    London   85      -5.0    190        220.0 -0.022727
4  2021-01-01  New York   60     -10.0    190        220.0 -0.045455
5  2021-01-01     Tokyo   45     -15.0    190        220.0 -0.068182
6  2022-01-01    London   70     -15.0    142        190.0 -0.078947
7  2022-01-01  New York   40     -20.0    142        190.0 -0.105263
8  2022-01-01     Tokyo   32     -13.0    142        190.0 -0.068421

优化后的链式调用实现 #

通过assign和分组操作的链式串联，避免频繁修改原DataFrame，提升代码可读性和简洁度：

result_df = (
    df
    # 计算每个日期的总Pop
    .assign(total=lambda x: x.groupby('Date')['Pop'].transform('sum'))
    # 按城市分组，依次计算差值、移位值及最终pc列
    .groupby('City', group_keys=False)
    .assign(
        pop_diff=lambda x: x['Pop'].diff(),
        total_shift=lambda x: x['total'].shift(),
        pc=lambda x: x['pop_diff'] / x['total_shift']
    )
)

print(result_df)

执行后输出结果与现有实现完全一致，但代码逻辑更连贯，无需创建中间变量或多次修改原数据。

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