如何通用化Pandas中销售数据9999缺失值的填充逻辑？

问题描述 #

处理销售数据时，需要替换DataFrame中值为9999的单元格。原始数据片段如下：

CUSTOMER_ID  ALL_Sales_2017  Toyota_sales_2017  Honda_sales_2017  Ford_sales_2017  9999_count
     3000522            93              9999                70                  20          1
     3000530            60              31                  9999                27          1
     3002817            231             9999                43                  170         1
     3004201            18              6                   9999                9999        2
     3004573            36              9999                18                  17          1
     3004888            9               9999                9999                9999        3

其中：

• ALL_Sales_YYYY列代表对应年份的总销量，无9999值
• 各品牌销量列（Toyota_sales_YYYY、Honda_sales_YYYY、Ford_sales_YYYY）可能存在1-3个9999值

填充逻辑：

• 单行含1个9999：总销量减去另外两个品牌的真实销量，结果填充到9999的位置
• 单行含2个9999：（总销量 - 唯一真实销量）除以2，平分填充到两个9999的位置
• 单行含3个9999：忽略该行，不做处理

处理后的数据示例：

CUSTOMER_ID  ALL_Sales_2017  Toyota_sales_2017  Honda_sales_2017  Ford_sales_2017  9999_count
    3000522            93              3                   70                  20          1
    3000530            60              31                  2                   27          1
    3002817            231             18                  43                  170         1
    3004201            18              6                   6                   6           2
    3004573            36              1                   18                  17          1
    3004888            9               9999                9999                9999        3

现有实现的局限 #

当前代码仅针对2017年硬编码处理，扩展性差，新增年份时需要重复编写相似逻辑：

df.loc[
    df["Toyota_sales_2017"].eq(9999) & (df["9999_count"] == 1), "Toyota_sales_2017"
] = (df["ALL_Sales_2017"] - df["Honda_sales_2017"] - df["Ford_sales_2017"])
df.loc[
    df["Honda_sales_2017"].eq(9999) & (df["9999_count"] == 1), "Honda_sales_2017"
] = (df["ALL_Sales_2017"] - df["Toyota_sales_2017"] - df["Ford_sales_2017"])
df.loc[df["Ford_sales_2017"].eq(9999) & (df["9999_count"] == 1), "Ford_sales_2017"] = (
    df["ALL_Sales_2017"] - df["Honda_sales_2017"] - df["Toyota_sales_2017"]
)

通用化解决方案 #

步骤1：提取所有年份 #

先从列名中提取所有存在的年份，避免硬编码：

# 提取所有年份（从ALL_Sales_YYYY列中获取）
years = [col.split("_")[-1] for col in df.columns if col.startswith("ALL_Sales_")]

步骤2：批量处理每个年份 #

对每个年份，动态获取对应的总销量列和品牌销量列，统一应用填充逻辑：

import pandas as pd

for year in years:
    # 定义当前年份的相关列
    total_col = f"ALL_Sales_{year}"
    brand_cols = [col for col in df.columns if col.endswith(f"_sales_{year}")]
    count_col = f"9999_count_{year}"  # 动态生成对应年份的count列
    
    # 计算当前年份每行的9999数量
    df[count_col] = df[brand_cols].eq(9999).sum(axis=1)
    
    # --- 处理1个9999的情况 ---
    for brand_col in brand_cols:
        # 获取其他品牌列
        other_brands = [col for col in brand_cols if col != brand_col]
        # 计算填充值：总销量 - 其他品牌销量之和
        fill_value = df[total_col] - df[other_brands].sum(axis=1)
        # 应用填充
        mask = df[brand_col].eq(9999) & (df[count_col] == 1)
        df.loc[mask, brand_col] = fill_value
    
    # --- 处理2个9999的情况 ---
    mask_2 = df[count_col] == 2
    # 计算每行的真实销量总和
    real_sales = df[brand_cols].where(df[brand_cols] != 9999).sum(axis=1)
    avg_fill = (df[total_col] - real_sales) / 2
    # 将两个9999的位置替换为平均值
    df.loc[mask_2, brand_cols] = df.loc[mask_2, brand_cols].replace(9999, avg_fill[mask_2].values[:, None])

代码说明 #

• 动态提取年份和对应列，新增年份无需修改代码
• 自动计算每行的9999数量，无需手动维护9999_count列
• 批量处理1个和2个9999的场景，避免重复编写每个品牌的逻辑
• 3个9999的情况自动忽略，无需额外处理

更简洁的优化方案 #

利用pandas向量化操作简化代码，减少循环嵌套：

import pandas as pd

# 提取所有年份
years = {col.split("_")[-1] for col in df.columns if col.startswith("ALL_Sales_")}

for year in years:
    total_col = f"ALL_Sales_{year}"
    brand_cols = [c for c in df.columns if c.endswith(f"_sales_{year}")]
    brand_df = df[brand_cols]
    
    # 计算每行9999的数量和真实销量总和
    count = brand_df.eq(9999).sum(axis=1)
    real_sum = brand_df.where(brand_df != 9999).sum(axis=1)
    
    # 处理1个9999的情况
    mask_1 = count == 1
    df.loc[mask_1, brand_cols] = df.loc[mask_1, brand_cols].mask(
        brand_df.eq(9999), df.loc[mask_1, total_col] - real_sum[mask_1], axis=0
    )
    
    # 处理2个9999的情况
    mask_2 = count == 2
    fill_val = (df.loc[mask_2, total_col] - real_sum[mask_2]) / 2
    df.loc[mask_2, brand_cols] = df.loc[mask_2, brand_cols].replace(9999, fill_val.values[:, None])

这个版本通过mask和replace的向量化操作，提升代码简洁性和执行效率。

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