在Pandas中计算同比差值并筛选各ID最大同比增长行

现有如下DataFrame：

ID YEAR     NPS
500 2020     0
500 2021     0
500 2022     0
501 2020    32      
501 2021    52      
501 2022    99      
503 2021    1       
503 2022    4       
504 2020    45      
504 2021    55      
504 2022    50      

需要完成两项操作：

1. 计算同比差值列nps_gain_yoy：每个ID的首次出现年份对应值为0，后续年份为当年NPS与上一年NPS的差值，得到中间结果。
2. 基于中间结果，为每个ID筛选出nps_gain_yoy最大的行，得到最终结果。

使用Pandas的分组、移位函数即可实现需求，具体步骤如下：

1. 构造原始DataFrame #

先将数据转换为Pandas DataFrame格式：

import pandas as pd

data = {
    'ID': [500,500,500,501,501,501,503,503,504,504,504],
    'YEAR': [2020,2021,2022,2020,2021,2022,2021,2022,2020,2021,2022],
    'NPS': [0,0,0,32,52,99,1,4,45,55,50]
}
df = pd.DataFrame(data)

2. 计算同比差值nps_gain_yoy #

按ID分组后，用shift()获取上一年的NPS值并计算差值，再将每组首行的空值填充为0：

# 分组计算当前NPS与上一年的差值
df['nps_gain_yoy'] = df.groupby('ID')['NPS'].diff()
# 填充每组首行的空值为0，并转为整数类型
df['nps_gain_yoy'] = df['nps_gain_yoy'].fillna(0).astype(int)

得到的中间结果如下：

ID  YEAR  NPS  nps_gain_yoy
0   500  2020    0             0
1   500  2021    0             0
2   500  2022    0             0
3   501  2020   32             0
4   501  2021   52            20
5   501  2022   99            47
6   503  2021    1             0
7   503  2022    4             3
8   504  2020   45             0
9   504  2021   55            10
10  504  2022   50            -5

3. 筛选每个ID的最大同比增长行 #

按ID分组后，通过idxmax()获取每组中nps_gain_yoy最大的行索引，再提取对应行数据：

# 获取每组中nps_gain_yoy最大的行
result = df.loc[df.groupby('ID')['nps_gain_yoy'].idxmax()]
# 重置索引（可选操作）
result = result.reset_index(drop=True)

最终结果如下：

ID  YEAR  NPS  nps_gain_yoy
0   500  2020    0             0
1   501  2022   99            47
2   503  2022    4             3
3   504  2021   55            10

（注：原示例中ID500的结果存在笔误，按规则其所有nps_gain_yoy均为0，因此最大行可取任意首次出现的行，此处取2020年数据）

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