含条件判断的Pandas嵌套循环未按预期运行问题排查

问题说明 #

需按以下规则为DataFrame生成新字段new_amt：

• 同一val_id分组内：
  1. 若product为CL，先取val_against与our_val_amt的较小值；若该分组内所有CL类的new_amt总和超过分组的our_val_amt，则调整CL类的new_amt，确保总和不超过额度（按顺序分配至额度用尽）
  2. 若product不为CL，new_amt取分组的our_val_amt减去CL类new_amt的总和
• 每个val_id独立执行上述逻辑

现有代码运行时抛出错误：Truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all()

测试数据集 #

import pandas as pd

data = {
    'val_id': ['V1', 'V1', 'V1', 'V2', 'V2'],
    'fac_id': ['F1', 'F2', 'F3', 'F4', 'F5'],
    'product': ['CL', 'CL', 'NON-CL', 'CL', 'NON-CL'],
    'val_against': [500, 600, 0, 300, 0],
    'our_val_amt': [1000, 1000, 1000, 400, 400]
}
df = pd.DataFrame(data)

现有错误代码 #

for val in df['val_id'].unique():
    group = df[df['val_id'] == val]
    total_cl = 0
    for idx, row in group.iterrows():
        if row['product'] == 'CL':
            temp = min(row['val_against'], row['our_val_amt'])
            if total_cl + temp > row['our_val_amt']:
                df.loc[idx, 'new_amt'] = row['our_val_amt'] - total_cl
                total_cl = row['our_val_amt']
            else:
                df.loc[idx, 'new_amt'] = temp
                total_cl += temp
        else:
            df.loc[idx, 'new_amt'] = row['our_val_amt'] - total_cl

预期输出 #

错误原因 #

报错源于循环中对Series的模糊布尔判断，且直接修改原DataFrame易引发索引匹配问题；另外原代码未统一取分组的our_val_amt额度，逻辑存在隐患。

正确解决方案 #

使用Pandas分组apply方法，对每个val_id分组单独处理，避免循环修改原DataFrame：

def process_group(group):
    # 获取分组统一额度
    total_limit = group['our_val_amt'].iloc[0]
    # 处理CL类数据
    cl_mask = group['product'] == 'CL'
    cl_rows = group[cl_mask].copy()
    # 初始计算CL候选值
    cl_rows['candidate'] = cl_rows[['val_against', 'our_val_amt']].min(axis=1)
    # 计算累计和判断是否超额度
    cl_rows['cum_sum'] = cl_rows['candidate'].cumsum()
    over_limit_idx = cl_rows[cl_rows['cum_sum'] > total_limit].index
    
    if not over_limit_idx.empty:
        first_over_idx = over_limit_idx[0]
        # 前序CL保持候选值
        cl_rows.loc[:first_over_idx-1, 'new_amt'] = cl_rows.loc[:first_over_idx-1, 'candidate']
        # 第一个超额度的CL分配剩余额度
        remaining = total_limit - cl_rows.loc[:first_over_idx-1, 'candidate'].sum()
        cl_rows.loc[first_over_idx, 'new_amt'] = remaining
        # 后续CL分配0
        cl_rows.loc[first_over_idx+1:, 'new_amt'] = 0
    else:
        cl_rows['new_amt'] = cl_rows['candidate']
    
    # 处理非CL类数据
    non_cl_rows = group[~cl_mask].copy()
    total_cl_amt = cl_rows['new_amt'].sum()
    non_cl_rows['new_amt'] = total_limit - total_cl_amt
    
    # 合并并还原原索引顺序
    processed = pd.concat([cl_rows, non_cl_rows]).reindex(group.index)
    return processed['new_amt']

# 应用分组处理
df['new_amt'] = df.groupby('val_id', group_keys=False).apply(process_group)

运行后即可得到预期输出。

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