如何在Pandas DataFrame中按ID分组逐行计算Qty_2减Qty_1并更新后续行

前置准备：构造测试数据 #

import pandas as pd

# 构造和示例一致的DataFrame
df = pd.DataFrame({
    'ID': ['A', 'A', 'A', 'B', 'B', 'B'],
    'Qty_1': [1, 2, 3, 3, 2, 1],
    'Qty_2': [10, 0, 0, 29, 0, 0]
})

需求1：分组首行保留原始Qty_2，后续行用上一行Qty_2减当前行Qty_1更新 #

自定义函数实现 #

def calc_v1(group):
    # 取当前ID组首行的原始Qty_2作为初始值
    init_qty2 = group['Qty_2'].iloc[0]
    # 生成扣减累计值：首行不扣减，后续行累计当前行的Qty_1
    cum_deduct = group['Qty_1'].mask(group.cumcount() == 0, 0).cumsum()
    group['Qty_2'] = init_qty2 - cum_deduct
    return group

df_result1 = df.groupby('ID', group_keys=False).apply(calc_v1)

一行简写实现 #

df['Qty_2'] = df.groupby('ID')['Qty_2'].transform('first') - df['Qty_1'].mask(df.groupby('ID').cumcount() == 0, 0).groupby(df['ID']).cumsum()

需求2：分组首行先扣减当前行Qty_1，后续行按相同规则更新 #

自定义函数实现 #

def calc_v2(group):
    # 取当前ID组首行的原始Qty_2作为初始值
    init_qty2 = group['Qty_2'].iloc[0]
    # 生成扣减累计值：从首行开始累计所有行的Qty_1
    cum_deduct = group['Qty_1'].cumsum()
    group['Qty_2'] = init_qty2 - cum_deduct
    return group

df_result2 = df.groupby('ID', group_keys=False).apply(calc_v2)

一行简写实现 #

df['Qty_2'] = df.groupby('ID')['Qty_2'].transform('first') - df.groupby('ID')['Qty_1'].cumsum()

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