如何按国家分组计算过去3年无缺失值的投资总和？

原始数据 #

需求说明 #

新增一列10 year investment，按国家分组计算过去3年的投资总和，仅当3年数据均无NaN时才计算，否则为NaN。以阿富汗为例，对应列值如下：

尝试过的错误代码 #

1. 带groupby的代码报错（移除groupby可运行）：

df['10 year investment'] = df.groupby('Country').sort_values(by='Year', ascending =False)['investment'].rolling(min_periods=3, window=3).sum()

2. 可运行但未按分组计算的代码：

df['10 year investment'] = df.sort_values(by='Year', ascending =False)['investment'].rolling(min_periods=3, window=3).sum()

正确实现方法 #

问题根源是groupby后直接调用sort_values的语法逻辑错误，正确的做法是先对整个DataFrame按国家+年份排序，再分组执行滚动计算：

# 第一步：按国家升序、年份降序排序，确保每个国家的年份从新到旧排列
df_sorted = df.sort_values(by=['Country', 'Year'], ascending=[True, False])

# 第二步：分组后对investment列做滚动求和，min_periods=3要求3个数据都非NaN才计算
df_sorted['10 year investment'] = df_sorted.groupby('Country')['investment'].rolling(window=3, min_periods=3).sum().reset_index(level=0, drop=True)

# 可选：如果需要恢复原DataFrame的原始顺序，执行以下代码
df = df_sorted.sort_index()

关键说明 #

• reset_index(level=0, drop=True)：移除分组产生的Country索引，让计算结果能和原DataFrame的行完美对齐
• min_periods=3：严格要求窗口内3个数据都不为NaN才返回求和结果，否则返回NaN，完全匹配需求
• 先整体排序再分组滚动：避免了groupby后直接排序的语法错误，确保每个国家的滚动计算独立进行

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