如何在Python的Pandas DataFrame中实现分组累计求和达标后重置

示例数据 #

假设你的原始DataFrame结构如下：

import pandas as pd

df = pd.DataFrame({
    'ac': ['A', 'A', 'A', 'A', 'B', 'B', 'B'],
    'Total_time': [100, 90, 50, 120, 200, 80, 150]
})

预期结果 #

需要生成的cum_total列需满足：按ac分组累计Total_time，当累计值超过185时，重置为当前行的Total_time并继续累计，最终结果如下：

问题分析 #

现有代码的核心问题是：重置累计时错误地将起始值设为0，而非当前行的Total_time值。正确逻辑应为：当当前累计值 + 当前行值超过阈值时，直接将当前行值作为新的累计起始值，而非清零。

修正代码 #

通过自定义分组处理函数实现需求：

def reset_cumsum(series, threshold=185):
    cum_values = []
    current_total = 0
    for val in series:
        # 判断累加后是否超过阈值
        if current_total + val > threshold:
            current_total = val
        else:
            current_total += val
        cum_values.append(current_total)
    return pd.Series(cum_values, index=series.index)

# 按ac分组应用函数
df['cum_total'] = df.groupby('ac')['Total_time'].apply(reset_cumsum)

代码说明 #

• 遍历分组内的每个Total_time值，维护current_total记录当前累计值
• 每次先判断累加当前值是否会超过阈值：
  • 超过则将current_total重置为当前值
  • 未超过则正常累加
• 最终返回分组内的累计结果序列，与原DataFrame索引对齐

错误代码对比（问题根源） #

如果之前的代码类似如下，就会出现重置后为0的问题：

# 错误示例
def wrong_reset(series):
    current = 0
    res = []
    for val in series:
        current += val
        if current > 185:
            current = 0  # 此处错误，应替换为val
        res.append(current)
    return pd.Series(res, index=series.index)

该代码在累计超阈值时直接清零，而非以当前行值作为新起点，这就是导致输出0的原因。

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