Pandas DataFrame条件处理：使用sum()累加并按规则删除行

处理规则 #

对给定的pandas DataFrame按如下逻辑处理：

• 按id字段分组遍历数据
• 若同一id对应多行且组内quantity存在负值，逐行累加quantity做核销处理
• 累加/核销后结果为0或负数的行直接删除，结果为正数的行保留

示例数据 #

输入df：

import pandas as pd
df = pd.DataFrame({
    'id': ['1a','1a','b5','b5','1a','1a'],
    'date':['11-01-22', '12-01-22', '13-01-22', '21-01-22', '11-01-22', '18-01-22'],
    'quantity':[2,5,3,-1,-2,2]
})

预期输出结果：

result = pd.DataFrame({
    'id': ['1a','b5','1a'],
    'date':['12-01-22', '13-01-22','18-01-22'],
    'quantity':[5,2,2]
})

实现方案 #

核心逻辑是先合并同id同日期的记录，再按时间顺序用负值逐笔核销之前的正值记录，核销完的负值记录直接丢弃，被核销为0的正值记录也丢弃，剩余的正值记录按日期排序即为结果：

# 按id、date分组聚合，合并同一天同id的正负数量
df_agg = df.groupby(['id', 'date'], as_index=False)['quantity'].sum()
# 过滤聚合后数量为0的记录，按时序排序
df_agg = df_agg[df_agg['quantity'] != 0].sort_values(['id', 'date']).reset_index(drop=True)

result_rows = []
# 按id分组处理负值核销
for _, group in df_agg.groupby('id', sort=False):
    group = group.sort_values('date').reset_index(drop=True)
    pos_buffer = []
    for _, row in group.iterrows():
        current_qty = row['quantity']
        if current_qty > 0:
            pos_buffer.append(row.copy())
        else:
            remain_neg = abs(current_qty)
            # 从最近的正值记录开始核销
            while remain_neg > 0 and pos_buffer:
                last_pos = pos_buffer[-1]
                if last_pos['quantity'] > remain_neg:
                    last_pos['quantity'] -= remain_neg
                    remain_neg = 0
                else:
                    remain_neg -= last_pos['quantity']
                    pos_buffer.pop()
    result_rows.extend(pos_buffer)

# 转换为DataFrame，按日期排序重置索引
result = pd.DataFrame(result_rows).sort_values('date').reset_index(drop=True)

执行代码后得到的结果与预期完全一致。

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