如何在Pandas DataFrame中按组逐行计算池额度消耗后的产品使用量

给定如下DataFrame（已修正原数据的语法错误）：

import pandas as pd

data = {'id':[1,1,1,2,2],
        'prod_type':['A','B','C','A','C'],
        'prod_order':[1,3,2,2,1],
        'prod_price':[100,150,300,10,40],
        'prod_pool':[200,200,200,45,45]
       }
df = pd.DataFrame(data)

需求说明：按id分组，每组内按prod_order排序后，逐行计算prod_used，规则如下：

• 每组首行：prod_used = min(prod_price, prod_pool)，剩余池额度 = 初始池额度 - prod_used
• 后续行：prod_used = min(prod_price, max(0, 上一步剩余池额度))，剩余池额度更新为max(0, 上一步剩余池额度 - prod_used)

解决方案 #

通过groupby结合自定义函数实现逐行动态计算：

def calculate_prod_used(group):
    # 按prod_order对组内数据排序
    group_sorted = group.sort_values('prod_order').reset_index(drop=True)
    # 初始化剩余池额度（取组内初始池额度）
    remaining_pool = group_sorted['prod_pool'].iloc[0]
    prod_used_list = []
    
    for price in group_sorted['prod_price']:
        # 计算当前行可使用的池额度
        used = min(price, remaining_pool)
        prod_used_list.append(used)
        # 更新剩余池额度，确保不小于0
        remaining_pool = max(0, remaining_pool - used)
    
    group_sorted['prod_used'] = prod_used_list
    # 计算每一步的剩余池额度（兜底修正为非负）
    group_sorted['remaining_pool'] = [group_sorted['prod_pool'].iloc[0] - sum(prod_used_list[:i+1]) for i in range(len(prod_used_list))]
    group_sorted['remaining_pool'] = group_sorted['remaining_pool'].clip(lower=0)
    return group_sorted

# 应用函数并合并结果
result_df = df.groupby('id').apply(calculate_prod_used).reset_index(drop=True)
print(result_df)

输出结果 #

运行代码后得到的最终结果：

id prod_type  prod_order  prod_price  prod_pool  prod_used  remaining_pool
0   1         A           1         100        200        100             100
1   1         C           2         300        200        100               0
2   1         B           3         150        200          0               0
3   2         C           1          40         45         40               5
4   2         A           2          10         45          5               0

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