Pandas按id、treatment分组实现累计和低于0时重置为0的方案

你原代码的问题在于：x.cumsum()返回的是分组内长度和分组一致的序列，不是单个标量，直接用if x.cumsum() < 0做判断会得到布尔序列，无法得到正确结果。

要实现累计和小于0就重置为0的逻辑，可以自定义累加函数，再作用到每个分组上，代码如下：

import pandas as pd

# 自定义带重置的累计和函数
def reset_cumsum(series):
    current_sum = 0
    result = []
    for val in series:
        current_sum += val
        # 累计和小于0就重置为0
        current_sum = max(current_sum, 0)
        result.append(current_sum)
    return pd.Series(result, index=series.index)

# 分组应用函数
df['cumsum'] = df.groupby(['id','treatment'])['value'].apply(reset_cumsum)

运行上述代码后得到的cumsum列就和你期望的输出完全匹配。

注：你示例输入里的数值列名为value，如果你的实际表中该列已改名为days，把代码里的'value'替换为'days'即可。

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