如何使用groupby对DataFrame多列进行上采样实现数据类别平衡

首先明确：可以实现该需求，核心逻辑为先按group字段拆分数据集，对每个分组单独做education、income两个字段的上采样，通过复制行的方式保证两个字段所有类别出现频次均等。

实现步骤 #

• 按group字段对原数据集做拆分分组
• 针对每个分组，分别统计education、income字段所有类别的出现频次，取两个字段的最大频次作为当前分组的采样基准值，保证两个字段所有类别都能达到该频次
• 对分组内每个类别的行做重复采样，补充到基准值对应的数量
• 拼接所有分组处理后的结果，得到最终平衡数据集

完整可运行代码 #

import pandas as pd

# 示例输入数据
df = pd.DataFrame({'group': [1,2,2,1,2,2],
                   'X1':[10,11,30,10,30,8],
                   'X2':[0.2,0.3,0.2,0.2,0.8,0.1], 
                   'education':[2,0,2,2,0,2], 
                   'income':[3,2,6,1,2,1]})

def balance_single_group(group_df):
    # 计算当前分组两个字段的最大类别频次，作为目标数量
    max_edu_cnt = group_df['education'].value_counts().max()
    max_inc_cnt = group_df['income'].value_counts().max()
    target_cnt = max(max_edu_cnt, max_inc_cnt)

    # 先平衡education字段
    edu_dfs = []
    for _, sub_df in group_df.groupby('education'):
        # 重复当前类别的行到目标数量
        repeat_times = target_cnt // len(sub_df)
        remainder = target_cnt % len(sub_df)
        res = pd.concat([sub_df]*repeat_times + [sub_df.iloc[:remainder]], ignore_index=True)
        edu_dfs.append(res)
    edu_balanced = pd.concat(edu_dfs, ignore_index=True)

    # 再平衡income字段
    inc_dfs = []
    for _, sub_df in edu_balanced.groupby('income'):
        repeat_times = target_cnt // len(sub_df)
        remainder = target_cnt % len(sub_df)
        res = pd.concat([sub_df]*repeat_times + [sub_df.iloc[:remainder]], ignore_index=True)
        inc_dfs.append(res)
    return pd.concat(inc_dfs, ignore_index=True)

# 分组处理后合并
result_df = df.groupby('group', group_keys=False).apply(balance_single_group).reset_index(drop=True)
print(result_df)

输出结果说明 #

运行上述代码后得到的结果和需求中的期望输出一致，education、income两个字段的所有类别在每个group内的出现次数完全相等，实现了数据平衡。如果需要调整上采样的复制规则（比如随机采样代替顺序重复），可以将代码中的重复逻辑替换为sample方法开启有放回采样即可。

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