如何编写Python函数基于阈值删除DataFrame中的分类变量列

原代码问题说明 #

• 你已经通过val_p.max()得到了单列最高占比的单值，不需要再加[i]索引，否则会直接报错
• var.append是列表的追加方法，不能直接赋值为df[var]，正确用法是追加符合保留条件的列名
• return缩进错误，放在了循环内部，导致只遍历第一列就提前返回结果
• 保留条件判断逻辑错误：仅当列的最高类别占比<=阈值时才保留该列，否则删除

正确代码实现 #

def dist_drop(df, cut):
    # 存储符合保留条件的列名
    keep_cols = []
    for col in df.columns:
        # 计算该列各分类的占比，取最高占比
        max_pct = df[col].value_counts(normalize=True).max() * 100
        # 最高占比不超过阈值则保留该列
        if max_pct <= cut:
            keep_cols.append(col)
    # 返回筛选后的DataFrame
    return df[keep_cols]

参数说明：df为待处理的全分类变量DataFrame，cut为百分比阈值（如阈值为51%直接传51即可）

测试验证 #

拿你给出的示例数据测试，结果和预期完全一致：

import pandas as pd
# 构造示例原始数据
data = {
    "Var1": ["Male", "Male", "Male", "Male"],
    "Var2": ["a", "b", "c", "c"],
    "Var3": ["Lite", "Full", "Full", "Lite"]
}
df = pd.DataFrame(data)
# 调用函数，传入阈值51%
output_df = dist_drop(df, 51)
print(output_df)

运行输出：

Var2  Var3
0    a  Lite
1    b  Full
2    c  Full
3    c  Lite

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