Pandas按月份升序逐行扣减还款额并标记还款状态

问题描述 #

现有存储抵押还款数据的DataFrame mortgage_data，包含name、mortgage_amount、month三列，其中month列已按升序排列。给定总还款额变量mortgage_amount_paid，需要按月份从小到大的顺序逐行抵扣mortgage_amount列的值，新增两列存储计算结果：

• mortgage_amount_updated：抵扣后的剩余应还金额
• paid_status：还款状态标记，规则如下：
  • 当月抵押金额被全额抵扣时，标记为full，mortgage_amount_updated为0
  • 还款总额不足以全额抵扣当月抵押金额时，标记为partial，mortgage_amount_updated为当月抵扣后的剩余金额
  • 还款总额已在之前月份耗尽、当月未被抵扣时，标记为zero，mortgage_amount_updated等于原mortgage_amount值

预期结果示例 #

当mortgage_amount_paid = 1000时，预期输出：

当mortgage_amount_paid = 600时，预期输出：

原有实现的边界错误 #

之前两种基于cumsum的实现均存在边界场景判断错误：

第一种错误实现 #

mortgage_amount_paid = 600

# 累计应还减总还款额
m1 = df['mortgage_amount'].cumsum().sub(mortgage_amount_paid)
# 判断累计应还是否超过总还款额
m2 = m1>0
# 判断上一行是否已经超过总还款额
m3 = m2.shift(fill_value=False)

df['mortgage_amount_updated'] = (m1.clip(0, mortgage_amount_paid)
                                   .mask(m3, df['mortgage_amount'])
                                 )
df['paid_status'] = np.select([m3, m2], ['zero', 'partial'], 'full')

错误表现：当mortgage_amount_paid=400时，预期paid_status应为full、zero、zero，实际输出为full、partial、zero。

第二种错误实现 #

mortgage_amount_paid = 600

m = df['mortgage_amount'].cumsum()

df['paid_status'] = np.select(
    [m <= mortgage_amount_paid,
     (m > mortgage_amount_paid) & (m.shift() < mortgage_amount_paid)
     ],
    ['full', 'partial'],
    default='zero'
)
df['mortgage_amount_updated'] = np.select(
    [df['paid_status'].eq('full'),
     df['paid_status'].eq('partial')],
    [0, m-mortgage_amount_paid],
    default=df['mortgage_amount']
)

错误表现：当mortgage_amount_paid=1时，预期paid_status应为partial、zero、zero，实际结果不符合预期。

正确实现代码 #

核心修复点为累计和偏移时首行填充0，避免NaN导致的条件判断失效，同时修正部分抵扣场景的剩余金额计算逻辑：

import numpy as np
import pandas as pd

def calc_mortgage_payment(df: pd.DataFrame, mortgage_amount_paid: float) -> pd.DataFrame:
    res = df.copy()
    # 计算累计应还金额
    cumsum_amt = res['mortgage_amount'].cumsum()
    # 上期累计应还，首行填充0解决边界判断问题
    prev_cumsum = cumsum_amt.shift(fill_value=0)

    # 标记还款状态
    res['paid_status'] = np.select(
        [
            cumsum_amt <= mortgage_amount_paid,
            (prev_cumsum < mortgage_amount_paid) & (cumsum_amt > mortgage_amount_paid)
        ],
        ['full', 'partial'],
        default='zero'
    )

    # 计算抵扣后剩余应还
    res['mortgage_amount_updated'] = np.select(
        [
            res['paid_status'].eq('full'),
            res['paid_status'].eq('partial')
        ],
        [
            0,
            res['mortgage_amount'] - (mortgage_amount_paid - prev_cumsum)
        ],
        default=res['mortgage_amount']
    )
    return res

边界场景验证 #

所有测试场景均符合预期：

• mortgage_amount_paid=400：输出状态为full/zero/zero，剩余应还为0/500/200
• mortgage_amount_paid=1：输出状态为partial/zero/zero，剩余应还为399/500/200
• mortgage_amount_paid=1100（覆盖全部应还总额）：输出状态为full/full/full，剩余应还全为0
• mortgage_amount_paid=0：输出状态为zero/zero/zero，剩余应还与原金额一致

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