Python编写函数实现分类变量基于Y=1占比的概率数值转换

实现逻辑 #

• 识别所有待转换分类列：自动排除因变量Y，其余列全部纳入转换范围
• 逐列分组统计：对每个分类列的每个类别，分别统计类别下总样本量、类别下Y取值为1的样本量
• 计算映射值：单类别映射值 = 类别下Y=1样本量 / 类别总样本量，自动处理无Y=1样本的类别（映射值为0）
• 批量映射替换：用计算好的映射值替换原分类取值，输出全数值型的新DataFrame
• 自动兼容Y的-1取值：统计Y=1样本数时会自动忽略取值为0、-1的样本，符合规则要求

通用函数代码 #

import pandas as pd
import numpy as np

def cat_to_prob_convert(input_df: pd.DataFrame, y_col: str = 'Y', keep_y: bool = True) -> pd.DataFrame:
    """
    批量将分类变量转换为对应类别下Y=1的样本占比
    参数:
        input_df: 原始数据集，需包含所有分类列与因变量列
        y_col: 因变量列名，默认值为'Y'，支持取值为-1/0/1
        keep_y: 结果中是否保留原始因变量Y列，默认保留
    返回:
        分类列完成概率转换后的数值型DataFrame
    """
    df = input_df.copy()
    cat_cols = [col for col in df.columns if col != y_col]
    res_df = pd.DataFrame()
    
    # 逐列计算映射关系
    for col in cat_cols:
        group_stats = df.groupby(col).agg(
            total = (y_col, 'count'),
            y1_count = (y_col, lambda x: (x == 1).sum())
        )
        val_map = (group_stats['y1_count'] / group_stats['total']).to_dict()
        res_df[f'{col}_convert'] = df[col].map(val_map)
    
    # 按参数判断是否保留Y列
    if keep_y:
        res_df[y_col] = df[y_col].values
    
    return res_df

示例运行验证 #

用提供的测试数据运行：

# 构造示例数据
df = pd.DataFrame([['iphone5', 'teams', 'shoe', 1],
['iphone6', 'teams', 'shirt', 0], ['iphone5', 'word', 'shoe', 0], 
['iphone7', 'ppt', 'pants', 0], ['iphone8', 'excel', 'umbrella', 1],
['iphone6', 'teams', 'shoe', 1], ['iphone9', 'publisher', 'food', 0]])
df.columns = ['Monday', 'Tuesday', 'Wednesday', 'Y']

# 执行转换
converted_result = cat_to_prob_convert(df, keep_y=True)
print(converted_result)

运行后输出的Monday_convert列取值和手动计算的结果完全一致：iphone5对应0.5、iphone6对应0.5、iphone7对应0、iphone8对应1、iphone9对应0，其余分类列也按相同规则完成转换。

对比手动写np.select的实现方式，该通用函数不需要提前枚举每个分类的所有取值，新增分类列、新增类别时不需要修改代码，也不会出现漏写类别导致的赋值错误，运行效率更高。

注意：该实现基于数据集内已出现的类别生成映射，如果后续有新类别进入数据集，需要重新用全量数据更新映射值。

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