如何仅对符合类别数量与占比要求的变量生成哑变量

问题描述 #

现有如下Pandas DataFrame，包含数值型与对象型列：

• 数据类型：
  • COL1 - 数值型
  • COL2 - 对象型
  • COL3 - 对象型

原始数据结构：

需要仅对满足以下两个条件的分类变量执行pandas.get_dummies()哑变量编码：

1. 变量的类别数量最多为3个
2. 变量中所有类别的占比均不低于0.4

规则示例 #

• COL2包含A、B、C、D共4个类别，不满足条件1，因此不进行编码
• COL3中类别"P"占比为1/5=0.2，不满足条件2，因此仅基于占比符合要求的"Y"和"Z"生成哑变量

期望输出 #

COL1 | COL3_Y | COL3_Z | ...  | COLn
-----|--------|--------|------|------
111  | 1      | 0      | ...  | ...
222  | 1      | 0      | ...  | ...
333  | 0      | 1      | ...  | ...
444  | 0      | 1      | ...  | ...
555  | 0      | 0      | ...  | ...

解决方案代码 #

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame({
    'COL1': [111, 222, 333, 444, 555],
    'COL2': ['A', 'A', 'B', 'C', 'D'],
    'COL3': ['Y', 'Y', 'Z', 'Z', 'P'],
    'COLn': ['x', 'x', 'y', 'y', 'z']
})

# 提取所有对象型分类列
categorical_cols = df.select_dtypes(include=['object']).columns.tolist()
cols_to_encode = []
valid_categories = {}

for col in categorical_cols:
    # 计算每个类别的占比
    cat_proportions = df[col].value_counts(normalize=True)
    # 检查条件：类别数≤3 且 所有类别占比≥0.4
    if len(cat_proportions) <= 3 and (cat_proportions >= 0.4).all():
        cols_to_encode.append(col)
        valid_categories[col] = cat_proportions.index.tolist()
    elif len(cat_proportions) > 3:
        # 类别数超3个，跳过
        continue
    else:
        # 仅保留占比≥0.4的类别
        valid_cats = cat_proportions[cat_proportions >= 0.4].index.tolist()
        if valid_cats:
            cols_to_encode.append(col)
            valid_categories[col] = valid_cats

# 生成指定类别的哑变量
dummies = pd.get_dummies(
    df[cols_to_encode],
    columns=cols_to_encode,
    prefix=cols_to_encode,
    prefix_sep='_',
    categories=valid_categories.values()
)

# 合并原始非编码列与哑变量列
result_df = pd.concat([df.drop(cols_to_encode, axis=1), dummies], axis=1)

print(result_df)

代码说明 #

1. 先筛选出所有对象型（分类）列，遍历每个列进行条件检查
2. 对每个分类列计算类别占比，判断是否满足两个编码条件
3. 对于部分类别符合占比要求的列，仅保留符合条件的类别用于编码
4. 使用pd.get_dummies()生成指定类别的哑变量，避免生成无效类别列
5. 最后将原始数据中的非编码列与生成的哑变量列合并，得到目标结果

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