Python/Pandas带负值重置为0的累积求和代码调试求助

需求说明 #

需要计算DataFrame中quantity列的累积和，规则为：当累积和为负值时重置为0，且需按product分组计算。当前代码运行结果与预期不符，需排查修复。

原始代码问题 #

原始代码尝试通过负值位置分组后计算累积和，逻辑如下：

neg = df['quantity'] < 0
df['py_output'] = df['quantity'].groupby([neg[::-1].cumsum(),df['product']]).cumsum().clip(0)

该逻辑错误在于：仅通过单个负值的位置反向分组，未考虑累积和跌破0的实际时机。例如连续多次卖出（负值）可能在累积和归0前仍有剩余额度，此时错误的分组会导致累积和计算偏差。

修正方案 #

由于此类带状态的累积计算依赖前一行结果，无法通过单纯的groupby.cumsum实现，需自定义累积函数并按分组应用：

1. 定义单个序列的累积计算函数：

def cumulative_sum_with_reset(series):
    current_sum = 0
    result = []
    for val in series:
        current_sum = max(current_sum + val, 0)
        result.append(current_sum)
    return pd.Series(result, index=series.index)

2. 按product分组应用函数：

df['py_output'] = df.groupby('product')['quantity'].apply(cumulative_sum_with_reset)

完整修正代码 #

import pandas as pd

data = [['Product-1', 'Time-1', '1. BUY', 1395, 1395]
        , ['Product-1', 'Time-2', '2. SELL', -9684, 0]
        , ['Product-1', 'Time-3', '1. BUY', 1352, 1352]
        , ['Product-1', 'Time-4', '2. SELL', -1348, 4]
        , ['Product-1', 'Time-5', '1. BUY', 1951, 1955]
        , ['Product-1', 'Time-6', '2. SELL', -1947, 8]
        , ['Product-1', 'Time-7', '1. BUY', 2554, 2562]
        , ['Product-1', 'Time-8', '1. BUY', 714, 3276]
        , ['Product-1', 'Time-9', '1. BUY', 445, 3721]
        , ['Product-1', 'Time-10', '1. BUY', 2948, 6669]
        , ['Product-1', 'Time-11', '1. BUY', 1995, 8664]
        , ['Product-1', 'Time-12', '2. SELL', -4161, 4503]
        , ['Product-1', 'Time-13', '2. SELL', -4161, 342]
        , ['Product-1', 'Time-14', '2. SELL', -2895, 0]
        , ['Product-1', 'Time-15', '1. BUY', 186, 186]
        , ['Product-1', 'Time-16', '1. BUY', 2646, 2832]
        , ['Product-1', 'Time-17', '1. BUY', 2594, 5426]
        , ['Product-1', 'Time-18', '2. SELL', -3202, 2224]
        , ['Product-1', 'Time-19', '1. BUY', 4170, 6394]
        , ['Product-1', 'Time-20', '1. BUY', 1766, 8160]
        , ['Product-1', 'Time-21', '2. SELL', -4403, 3757]
        , ['Product-1', 'Time-22', '2. SELL', -3523, 234]
        , ['Product-1', 'Time-23', '1. BUY', 1403, 1637]
        , ['Product-1', 'Time-24', '1. BUY', 1566, 3203]
        , ['Product-1', 'Time-25', '2. SELL', -1357, 1846]
        , ['Product-1', 'Time-26', '2. SELL', -1566, 280]
        , ['Product-1', 'Time-27', '1. BUY', 791, 1071]
        , ['Product-1', 'Time-28', '1. BUY', 2384, 3455]
        , ['Product-1', 'Time-29', '1. BUY', 1292, 4747]
        , ['Product-1', 'Time-30', '1. BUY', 1343, 6090]
        , ['Product-1', 'Time-31', '1. BUY', 322, 6412]
        , ['Product-2', 'Time-1', '1. BUY', 1248, 1248]
        , ['Product-2', 'Time-2', '1. BUY', 3276, 4524]
        , ['Product-2', 'Time-3', '1. BUY', 707, 5231]
        , ['Product-2', 'Time-4', '2. SELL', -3534, 1697]
        , ['Product-2', 'Time-5', '1. BUY', 1358, 3055]
        , ['Product-2', 'Time-6', '1. BUY', 253, 3308]
        , ['Product-2', 'Time-7', '2. SELL', -1082, 2226]
        , ['Product-2', 'Time-8', '1. BUY', 238, 2464]
        , ['Product-2', 'Time-9', '1. BUY', 371, 2835]]

cols = ['product', 'time', 'activity', 'quantity', 'desired_output']
 
df = pd.DataFrame(data, columns=cols)

# 修正后的计算逻辑
def cumulative_sum_with_reset(series):
    current_sum = 0
    result = []
    for val in series:
        current_sum = max(current_sum + val, 0)
        result.append(current_sum)
    return pd.Series(result, index=series.index)

df['py_output'] = df.groupby('product')['quantity'].apply(cumulative_sum_with_reset)

print(df)

运行后py_output将与desired_output完全匹配。

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