如何用Pandas按城市和ID计算aa_cumul累计列？代码出错求修正

需求说明 #

需要按city和ID分组创建aa_cumul列：

• 每组首行取new_r_aa与cml_aa_bx的和作为初始值
• 后续每行累加当前行的new_r_aa值

原始数据 #

import pandas as pd
data = {
    'city': ['NY', 'NY', 'NY', 'NY', 'NY', 'CA'],
    'ID': ['AA', 'AA', 'AA', 'AA', 'AA', 'AA'],
    'cml_aa_bx': [1, 3, 6, 10, 12, 2],
    'new_r_aa': [2, 6, 9, 8, 6, 5]
}

df = pd.DataFrame(data)

期望结果 #

data = {
    'city': ['NY', 'NY', 'NY', 'NY', 'NY', 'CA'],
    'ID': ['AA', 'AA', 'AA', 'AA', 'AA', 'AA'],
    'cml_aa_bx': [1, 3, 6, 10, 12, 2],
    'new_r_aa': [2, 6, 9, 8, 6, 5],
    'aa_cumul': [3, 11, 20, 28, 34, 6]
}

错误代码问题分析 #

当前代码未处理分组逻辑，直接从第一行开始全局累加，导致跨组（如NY到CA）时未重置初始值，同时未按组内规则计算，结果不符合要求。

修正后的代码 #

循环实现（直观对应需求逻辑） #

import pandas as pd

data = {
    'city': ['NY', 'NY', 'NY', 'NY', 'NY', 'CA'],
    'ID': ['AA', 'AA', 'AA', 'AA', 'AA', 'AA'],
    'cml_aa_bx': [1, 3, 6, 10, 12, 2],
    'new_r_aa': [2, 6, 9, 8, 6, 5]
}

df = pd.DataFrame(data)

# 按city和ID分组计算aa_cumul
def calculate_cumul(group):
    # 初始化累加序列，首行取new_r_aa与cml_aa_bx的和
    cumul = [group['new_r_aa'].iloc[0] + group['cml_aa_bx'].iloc[0]]
    # 遍历组内后续行，累加当前行的new_r_aa
    for i in range(1, len(group)):
        cumul.append(cumul[-1] + group['new_r_aa'].iloc[i])
    group['aa_cumul'] = cumul
    return group

df = df.groupby(['city', 'ID'], group_keys=False).apply(calculate_cumul)
print(df)

向量化实现（高效简洁） #

import pandas as pd

data = {
    'city': ['NY', 'NY', 'NY', 'NY', 'NY', 'CA'],
    'ID': ['AA', 'AA', 'AA', 'AA', 'AA', 'AA'],
    'cml_aa_bx': [1, 3, 6, 10, 12, 2],
    'new_r_aa': [2, 6, 9, 8, 6, 5]
}

df = pd.DataFrame(data)

# 利用cumsum实现组内累加，结合组首cml_aa_bx值得到结果
df['aa_cumul'] = df.groupby(['city', 'ID']).apply(
    lambda g: g['cml_aa_bx'].iloc[0] + g['new_r_aa'].cumsum()
).reset_index(drop=True)
print(df)

代码解释 #

1. 分组处理：通过groupby(['city', 'ID'])将数据按城市和ID拆分，确保每组独立计算累加值，避免跨组干扰
2. 初始值计算：每组首行取new_r_aa与cml_aa_bx的和作为累加起点
3. 后续行累加：循环实现中手动遍历组内后续行，基于前一行结果累加当前new_r_aa；向量化实现利用cumsum()直接计算new_r_aa的累加和，再加上组内第一个cml_aa_bx值，逻辑与循环完全一致但效率更高

注：按需求规则，CA组的aa_cumul应为2+5=7，与期望结果中的6存在差异，推测为期望数据笔误。

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