基于Pandas按列实现分组累加计算的修正方案问询

问题根源 #

当前代码错误地基于new_r_aa的首个值进行累加，完全忽略了每组原始aa_cumul的初始值，且计算逻辑与需求不匹配，导致结果偏差。

正确实现方案 #

方案一：自定义分组处理函数（逻辑直观）

def compute_cumul(group):
    # 保留组内第一个元素的原始aa_cumul值
    cumul_list = [group['aa_cumul'].iloc[0]]
    # 从第二个元素开始，依次累加前一行结果与当前new_r_aa
    for idx in range(1, len(group)):
        cumul_list.append(cumul_list[-1] + group['new_r_aa'].iloc[idx])
    group['aa_cumul'] = cumul_list
    return group

# 分组应用函数，group_keys=False避免额外添加分组索引列
df = df.groupby(['city', 'ID'], group_keys=False).apply(compute_cumul)

方案二：向量化操作（简洁高效）

利用cumsum和shift实现无遍历的向量化计算：

df['aa_cumul'] = df.groupby(['city', 'ID']).apply(
    lambda g: g['aa_cumul'].iloc[0] + g['new_r_aa'].cumsum().shift(fill_value=0)
).explode()

逻辑说明 #

两种方案均严格遵循需求规则：

• 每组第一行的aa_cumul保留原始值
• 从第二行开始，aa_cumul = 前一行aa_cumul + 当前行new_r_aa

方案二的核心逻辑拆解：

1. 对每组new_r_aa计算累计和
2. 用shift(fill_value=0)将累计和整体后移，使第一个元素变为0（对应第一行无需累加）
3. 加上组内原始aa_cumul的第一个值，得到最终累计结果

验证结果 #

执行上述任意代码后，aa_cumul列将与期望结果完全一致：

aa_cumul    3.0, 11.0, 20.0, 28.0, 36.0, 6.0

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