Python Pandas：基于累计和扣除DataFrame金额，如何生成amount_extra列？

数据准备 #

现有两个DataFrame：

import pandas as pd

df1 = pd.DataFrame({
    'name': ['virat', 'virat', 'virat', 'virat', 'virat', 'rohit', 'rohit'],
    'amount': [1000, 500, 200, 500, 150, 120, 100]})

df2 = pd.DataFrame({
    'name': ['virat', 'rohit'],
    'amount_paid': [1400, 300]})

需求说明 #

df2的name字段唯一，需使用df2的amount_paid列，按顺序扣除df1中对应name的amount值，直至amount_paid耗尽；若存在超额支付金额，新增amount_extra列记录该值（仅在对应name的最后一行显示超额金额，其余行显示0）。

预期输出 #

name    amount_left   amount_used_to_knock_off      amount_extra
0   virat   0                1000                          0
1   virat   100               400                          0
2   virat   200                 0                          0
3   virat   500                 0                          0
4   virat   150                 0                          0
5   rohit   0                 120                          0
6   rohit   0                 100                         80

已实现的代码及当前输出 #

以下代码已成功生成amount_left和amount_used_to_knock_off列，但缺少amount_extra列：

s = (df1
 .groupby('name')['amount'].cumsum()
 .sub(df1['name'].map(df2.set_index('name')['amount_paid']))
)

df1['amount_left'] = df1['amount'].mask(s.lt(df1['amount']), s.clip(lower=0))

df1['amount_used_to_knock_off'] = (df1['amount'].sub(s).clip(lower=0)
                                   .mask(s<0, df1['amount'])
                                   )

当前输出：

name  amount  amount_left  amount_used_to_knock_off
0  virat    1000            0                      1000
1  virat     500          100                       400
2  virat     200          200                         0
3  virat     500          500                         0
4  virat     150          150                         0
5  rohit     120            0                       120
6  rohit     100            0                       100

解决方案：新增amount_extra列 #

在现有代码基础上，添加以下逻辑计算超额支付金额：

简洁实现版 #

# 计算每个name的总金额
total_amount_per_name = df1.groupby('name')['amount'].transform('sum')
# 获取每个name对应的已支付金额
paid_per_name = df1['name'].map(df2.set_index('name')['amount_paid'])

# 初始化超额列
df1['amount_extra'] = 0
# 定位每个name的最后一行，计算并赋值超额金额
last_rows = df1.groupby('name').tail(1).index
df1.loc[last_rows, 'amount_extra'] = (paid_per_name - total_amount_per_name).loc[last_rows].clip(lower=0)

执行后即可得到符合预期的输出。

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