如何修正Pandas自定义函数以正确计算new_field列

需求规则 #

• 同一val_id分组内：
  1. 若product为CL：new_field取val_against与our_val_amt的最小值；若CL的new_field累计和超过our_val_amt，超出行的new_field取our_val_amt减去之前的累计和（如val_id=xx4时，200+300+50=550>510，最后一行new_field为510-500=10）。
  2. 若product不为CL：new_field取our_val_amt减去该分组内CL的new_field总和；若分组内无CL产品，则将our_val_amt按行分配（如val_id=xx7时，700先分配650到第一行，剩余50到第二行，后续行取0）。
  3. 每个val_id分组独立执行上述逻辑。

现有代码问题 #

当前代码仅能正确处理name=compx的分组（0-9行），后续compy、compz分组结果不符合预期：

• compy的xx5分组中，DL行的new_field应为0，实际输出为10
• compz的xx7分组中，无CL产品，应按行分配700（650、50、0），实际输出全为700
• compz的xx6分组中，DL行的new_field应为0，实际输出为50

现有代码执行逻辑解析 #

1. compute_new_field_for_cl函数 #

该函数用于处理CL行的new_field计算：

• 将tuple列拆分为DataFrame，给CL行赋值min(our_val_amt, val_against)，非CL行赋值0
• 计算new_field的累计和cumsum
• 将累计和超过our_val_amt的行的new_field设为0，再计算修正后的累计和max_cumsum
• 对累计和超标的行，用our_val_amt - max_cumsum得到剩余额度
• 问题：处理超出行的逻辑存在漏洞，且未考虑后续非CL行的关联处理

2. compute_new_field_for_not_cl函数 #

该函数用于处理非CL行的new_field计算：

• 将tuple列拆分为DataFrame，计算非CL行的our_val_amt减去CL行的new_field总和
• 填充NaN为0后加上原new_field（保留CL行的数值）
• 问题：
  • 未处理分组内无CL的场景，导致这类分组直接返回our_val_amt
  • 计算非CL行剩余额度时，索引匹配逻辑错误，导致部分行结果异常

修正后的完整代码 #

import pandas as pd

# 初始化原始数据
df = pd.DataFrame(data=[["compx","xx1","yy1",424,418,"XL"],["compx","xx1","yy2",424,134,"CL"],["compx","xx2","yy3",472,60,"DL"],["compx","xx2","yy4",472,104,"CL"], ["compx", "xx3", "yy5", 490, 50, "XL"], ["compx", "xx3", "yy6", 490, 500, "CL"], ["compx", "xx3", "yy7", 490, 200, "DL"], ["compx", "xx4", "yy8", 510, 200, "CL"], ["compx", "xx4", "yy9", 510, 300, "CL"], ["compx", "xx4", "yy10", 510, 50, "CL"], ["compy", "xx5", "yy11", 510, 200, "CL"], ["compy", "xx5", "yy12", 510, 300, "CL"], ["compy", "xx5", "yy12", 510, 50, "CL"], ["compy", "xx5", "yy13", 510, 30, "DL"], ["compz", "xx6", "yy14", 350, 200, "CL"], ["compz", "xx6", "yy15", 350, 100, "CL"], ["compz", "xx6", "yy16", 350, 50, "XL"], ["compz", "xx6", "yy17", 350, 50, "DL"], ["compz", "xx7", "yy18", 700, 650, "DL"], ["compz", "xx7", "yy19", 700, 200, "DL"], ["compz", "xx7", "yy20", 700, 400, "XL"] ], columns=["name","val_id","fac_id","our_val_amt","val_against","product"])

def calculate_new_field(group):
    # 获取分组的our_val_amt（同一分组内值相同）
    total_amt = group['our_val_amt'].iloc[0]
    group = group.copy()
    group['new_field'] = 0
    
    # 处理CL行
    cl_mask = group['product'] == 'CL'
    if cl_mask.any():
        # 初始赋值min(val_against, total_amt)
        group.loc[cl_mask, 'new_field'] = group.loc[cl_mask, 'val_against'].apply(lambda x: min(x, total_amt))
        # 计算累计和
        group['cumsum_cl'] = group['new_field'].cumsum()
        # 找出累计和超过total_amt的位置
        over_mask = group['cumsum_cl'] > total_amt
        if over_mask.any():
            # 找到第一个超标的行
            first_over_idx = over_mask.idxmax()
            # 前面的行保持不变，超出行的new_field调整为剩余额度
            remaining = total_amt - group.loc[:first_over_idx, 'new_field'].sum() + group.loc[first_over_idx, 'new_field']
            group.loc[first_over_idx, 'new_field'] = remaining if remaining > 0 else 0
            # 超出行之后的CL行设为0
            group.loc[group.index > first_over_idx, 'new_field'] = 0
        # 计算CL的总分配额
        cl_total = group['new_field'].sum()
        # 处理非CL行
        non_cl_mask = ~cl_mask
        remaining_non_cl = total_amt - cl_total
        if remaining_non_cl > 0:
            group.loc[non_cl_mask, 'new_field'] = 0
            # 按顺序分配剩余额度
            current_remaining = remaining_non_cl
            for idx in group[non_cl_mask].index:
                if current_remaining <= 0:
                    break
                assign = min(group.loc[idx, 'val_against'], current_remaining)
                group.loc[idx, 'new_field'] = assign
                current_remaining -= assign
        else:
            group.loc[non_cl_mask, 'new_field'] = 0
    else:
        # 无CL产品，按行分配total_amt
        current_remaining = total_amt
        for idx in group.index:
            if current_remaining <= 0:
                group.loc[idx, 'new_field'] = 0
                continue
            assign = min(group.loc[idx, 'val_against'], current_remaining)
            group.loc[idx, 'new_field'] = assign
            current_remaining -= assign
    # 清理临时列
    if 'cumsum_cl' in group.columns:
        group = group.drop('cumsum_cl', axis=1)
    return group

# 应用函数到每个val_id分组
df = df.groupby('val_id', group_keys=False).apply(calculate_new_field)
print(df)

修正说明 #

1. 改用分组整体处理：直接对val_id分组的整个DataFrame进行操作，避免了原代码中用tuple传递列的复杂逻辑，可读性和可维护性更强。
2. 优化CL行计算：
   • 先给CL行赋初始值，再计算累计和，精准定位超出行并调整剩余额度，确保CL行的总和不超过our_val_amt。
3. 完善非CL行处理：
   • 计算CL行总分配额后，若有剩余额度，按行顺序分配给非CL行；若无剩余则非CL行全为0。
   • 单独处理无CL产品的分组，按行顺序分配our_val_amt，直到额度耗尽。
4. 确保数据一致性：每个分组的new_field总和严格等于our_val_amt，符合规则要求。

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