Python如何实现支持二分类多分类的数据不平衡预检测函数？

我们可以基于最大类别样本数与最小类别样本数的比值作为判断依据，该逻辑天然适配任意类别数的二分类、多分类场景，不需要提前指定类别数量。默认阈值取业界常用的4（即最大类样本量是最小类的4倍以上时，判定为需要做平衡处理），阈值支持根据业务需求自定义调整。

函数实现代码 #

import numpy as np
from collections import Counter

def check_balance(y, imbalance_threshold=4.0) -> bool:
    """
    检测数据集类别分布是否需要做平衡处理
    Args:
        y: 类别标签数组，支持列表、numpy数组、pandas Series格式
        imbalance_threshold: 不平衡判定阈值，最大类样本数/最小类样本数超过该值则返回True（需要平衡）
    Returns:
        bool: True表示存在类别不平衡需要做平衡处理，False表示分布均衡
    """
    # 统计每个类别的样本数
    label_counts = Counter(y)
    # 少于2个类别时无分类任务意义，直接返回不需要平衡
    if len(label_counts) < 2:
        return False
    # 获取最大、最小类别的样本量
    max_count = max(label_counts.values())
    min_count = min(label_counts.values())
    # 按阈值判断是否需要平衡
    return (max_count / min_count) > imbalance_threshold

使用示例 #

• 二分类不平衡场景测试

y = [0]*900 + [1]*100
print(check_balance(y)) 
# 输出：True （比例9:1超过默认阈值4，判定为需要平衡）

• 多分类均衡场景测试

y = [0]*300 + [1]*300 + [2]*300
print(check_balance(y))
# 输出：False （最大最小样本量比值为1，分布均衡）

• 自定义阈值场景

y = [0]*800 + [1]*200
print(check_balance(y, imbalance_threshold=5))
# 输出：False （比例4:1小于自定义阈值5，符合均衡要求）

你可以根据业务场景调整imbalance_threshold的取值：对类别不平衡敏感度高的场景可以设为2.0，普通分类场景默认4.0即可，对分类效果要求宽松的场景可以设为10.0。

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