Pandas按国家分组计算id_score与下一行Dem_Score的差值

问题背景 #

原始DataFrame如下：

countryname  yr  US_Election_Year   id_score       Dem_Score
  Albania   1992    1990            0.688809        0.366570                        
  Albania   1997    1996            0.024751        0.247750            
  Argentina 1995    1992            0.081818        0.398908            
  Argentina 1999    1996            -0.521796       0.247759
  Argentina 2003    2000            -0.293386       -0.102298

需要实现：按countryname分组，计算每组内第i行的id_score减去第i+1行的Dem_Score，最终生成Delta_Dem_Dist列。例如阿根廷1999年的行，差值为0.081818 - 0.247759 = -0.165941，期望结果如下：

countryname yr  US_Election_Year    id_score    Dem_Score   Delta_Dem_Dist
  Albania   1992    1990            0.688809    0.366570                            
  Albania   1997    1996            0.024751    0.247750        0.441130
  Argentina 1995    1992            0.081818    0.398908                    
  Argentina 1999    1996            -0.521796   0.247759        -0.165941   
  Argentina 2003    2000            -0.293386   -0.102298       -0.624094   

直接使用df['id_score'] - df['Dem_Score'].shift(-1)会跨分组计算，无法满足需求，需要针对每组单独处理。

解决方案 #

利用Pandas的groupby结合shift方法，在每个国家分组内完成移位和差值计算：

import pandas as pd

# 构造原始数据（如果已有DataFrame可跳过此步骤）
data = {
    'countryname': ['Albania', 'Albania', 'Argentina', 'Argentina', 'Argentina'],
    'yr': [1992, 1997, 1995, 1999, 2003],
    'US_Election_Year': [1990, 1996, 1992, 1996, 2000],
    'id_score': [0.688809, 0.024751, 0.081818, -0.521796, -0.293386],
    'Dem_Score': [0.366570, 0.247750, 0.398908, 0.247759, -0.102298]
}
df = pd.DataFrame(data)

# 核心计算逻辑：按国家分组后对Dem_Score移位，再计算差值
df['Delta_Dem_Dist'] = df['id_score'] - df.groupby('countryname')['Dem_Score'].shift(-1)

代码解释 #

• df.groupby('countryname')['Dem_Score'].shift(-1)：对每个国家的Dem_Score列执行向上移位（取当前行的下一行值），确保移位操作仅在同一国家组内进行，不会跨组干扰。
• df['id_score'] - ...：用当前行的id_score减去移位后的Dem_Score，得到目标差值。每个国家的最后一行因无下一行数据，结果会显示为NaN，与期望输出一致。

运行上述代码后，即可得到符合要求的DataFrame。

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