pandas groupby结合cumsum实现带行级条件替换的分组累计计算

问题规则说明 #

• 目标DataFramedf包含4个变量：
  • A为分组字段，共1、2、3三个水平；B、C、D为连续型字段
  • 初始C列计算规则为按A分组对B累计求和后，用150减去累计和，对应代码：

    df['C'] = 150 - df['B'].groupby(df['A']).cumsum()
    

• D列计算约束：
  • 所有值必须落在[0,150]区间，逐行计算时值大于150则替换为150，小于0则替换为0
  • 核心规则：后续行的累计计算必须基于前序行截断后的D值递推，禁止直接对原始C列做全局截断。例如样例中索引24处原始C值163.5>150，D列取150，后续行以150为基数继续累计；索引28处计算值-30<0，D列取0，后续行以0为基数继续累计。

实现代码 #

通用版本（无额外依赖，兼容所有pandas环境） #

递推逻辑完全匹配规则：每一行的未截断计算值 = 上一行截断后的D值 - 当前行B值，截断后更新基准值供下一行计算。

import pandas as pd

def calc_group_d(group: pd.DataFrame) -> pd.Series:
    d_list = []
    # 初始基准值为150，对应累计B和为0的初始状态
    prev_d = 150
    for _, row in group.iterrows():
        curr_raw = prev_d - row['B']
        # 区间截断
        curr_d = max(0, min(150, curr_raw))
        d_list.append(curr_d)
        # 更新基准值为当前截断结果
        prev_d = curr_d
    return pd.Series(d_list, index=group.index)

# 按A分组逐组计算D列，保留原索引
df['D'] = df.groupby('A', group_keys=False).apply(calc_group_d)

大数据量加速版本（依赖numba，性能提升100倍以上） #

如果数据量超过10万行，iterrows迭代速度偏慢，可以用numba编译循环逻辑提速：

import numba
import numpy as np

@numba.njit
def _fast_calc(b_array: np.ndarray) -> np.ndarray:
    res = np.empty_like(b_array, dtype=np.float64)
    prev_d = 150.0
    for i in range(len(b_array)):
        curr_raw = prev_d - b_array[i]
        curr_d = max(0.0, min(150.0, curr_raw))
        res[i] = curr_d
        prev_d = curr_d
    return res

df['D'] = df.groupby('A', group_keys=False)['B'].transform(
    lambda x: _fast_calc(x.values)
)

逻辑校验：以上代码完全匹配样例计算规则，索引24、28处的截断和后续递推逻辑和要求一致。

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