Pandas下采样时如何保留各列按target分组的频率分布？

需求说明 #

我有一个示例pandas dataframe df，数据结构如下：

{'column_a': {0: 'b',   1: 'b',  2: 'a',  3: 'b',  4: 'd',  5: 'a',  6: 'b',  7: 'b',  8: 'c',  9: 'a',  10: 'a',  11: 'a',  12: 'a',  13: 'c',  14: 'c',  15: 'c',  16: 'b',  17: 'a',  18: 'a',  19: 'b',  20: 'd',  21: 'c',  22: 'a',  23: 'b',  24: 'c',  25: 'c',  26: 'c',  27: 'e',  28: 'e',  29: 'e',  30: 'e',  31: 'c',  32: 'e',  33: 'e',  34: 'd',  35: 'e',  36: 'd',  37: 'e',  38: 'd',  39: 'b',  40: 'd',  41: 'c',  42: 'b',  43: 'd',  44: 'c',  45: 'e',  46: 'd',  47: 'c',  48: 'e',  49: 'b',  50: 'c'}, 'column_b': {0: 'c',  1: 'b',  2: 'b',  3: 'd',  4: 'b',  5: 'a',  6: 'd',  7: 'c',  8: 'c',  9: 'd',  10: 'a',  11: 'a',  12: 'b',  13: 'a',  14: 'c',  15: 'd',  16: 'd',  17: 'c',  18: 'b',  19: 'd',  20: 'a',  21: 'a',  22: 'd',  23: 'b',  24: 'a',  25: 'c',  26: 'e',  27: 'd',  28: 'b',  29: 'c',  30: 'd',  31: 'b',  32: 'e',  33: 'b',  34: 'b',  35: 'c',  36: 'b',  37: 'b',  38: 'd',  39: 'c',  40: 'b',  41: 'a',  42: 'b',  43: 'e',  44: 'e',  45: 'c',  46: 'e',  47: 'c',  48: 'b',  49: 'b',  50: 'c'}, 'column_c': {0: 'b',  1: 'd',  2: 'b',  3: 'b',  4: 'd',  5: 'c',  6: 'b',  7: 'a',  8: 'a',  9: 'a',  10: 'a',  11: 'b',  12: 'd',  13: 'c',  14: 'b',  15: 'a',  16: 'a',  17: 'a',  18: 'b',  19: 'c',  20: 'a',  21: 'a',  22: 'b',  23: 'd',  24: 'd',  25: 'c',  26: 'd',  27: 'c',  28: 'c',  29: 'e',  30: 'd',  31: 'c',  32: 'd',  33: 'c',  34: 'b',  35: 'b',  36: 'd',  37: 'd',  38: 'd',  39: 'b',  40: 'c',  41: 'e',  42: 'e',  43: 'b',  44: 'b',  45: 'd',  46: 'd',  47: 'c',  48: 'e',  49: 'd',  50: 'b'}, 'column_d': {0: 'b',  1: 'c',  2: 'd',  3: 'd',  4: 'b',  5: 'b',  6: 'd',  7: 'd',  8: 'd',  9: 'b',  10: 'd',  11: 'c',  12: 'b',  13: 'a',  14: 'c',  15: 'c',  16: 'd',  17: 'c',  18: 'd',  19: 'a',  20: 'd',  21: 'b',  22: 'd',  23: 'b',  24: 'd',  25: 'e',  26: 'c',  27: 'c',  28: 'c',  29: 'd',  30: 'c',  31: 'e',  32: 'd',  33: 'd',  34: 'd',  35: 'b',  36: 'c',  37: 'e',  38: 'b',  39: 'e',  40: 'b',  41: 'c',  42: 'b',  43: 'e',  44: 'b',  45: 'c',  46: 'd',  47: 'c',  48: 'c',  49: 'b',  50: 'd'}, 'target': {0: 1,  1: 1,  2: 1,  3: 1,  4: 1,  5: 1,  6: 1,  7: 1,  8: 1,  9: 1,  10: 1,  11: 1,  12: 1,  13: 1,  14: 1,  15: 1,  16: 1,  17: 1,  18: 1,  19: 1,  20: 1,  21: 1,  22: 1,  23: 1,  24: 1,  25: 0,  26: 0,  27: 0,  28: 0,  29: 0,  30: 0,  31: 0,  32: 0,  33: 0,  34: 0,  35: 0,  36: 0,  37: 0,  38: 0,  39: 0,  40: 0,  41: 0,  42: 0,  43: 0,  44: 0,  45: 0,  46: 0,  47: 0,  48: 0,  49: 0,  50: 0}}

需要对该DataFrame按任意指定行数或比例下采样，同时尽可能保留每个target类别下各列不同取值的频率分布。
直接使用.sample()方法执行下采样，如smaller_df = df.sample(n=100)或smaller_df = df.sample(frac=0.1)，很可能破坏每个类别下各列的取值分布。
可以通过如下代码查看当前的频率密度：

for col in df.columns:
    print(df.groupby(['target'])[col].value_counts(normalize=True))

输出示例如下：

target        column_a
0       e           0.384615
        c           0.269231
        d           0.230769
        b           0.115385
1       a           0.360000
        b           0.320000
        c           0.240000
        d           0.080000

现有公开方案仅能解决保留单一分布的下采样问题，无法满足多列分布同时保留的需求，真实数据集规模为(8370994, 731)。

实现方案 #

采用分层分组抽样逻辑，先按target列拆分不同类别子集，每个子集内独立按比例/指定行数随机抽样，既能保留target的原始分布，也能保证每个target分组内所有列的频率分布和原数据完全一致。

按比例下采样代码 #

import pandas as pd

# 自定义下采样比例，例如下采样到原数据的10%
sample_frac = 0.1
# random_state固定可保证结果可复现
sampled_df = df.groupby('target', group_keys=False).apply(
    lambda x: x.sample(frac=sample_frac, random_state=42)
)

按指定总行数下采样代码 #

import pandas as pd

# 自定义总采样行数，例如采样10000行
total_sample_n = 10000
# 计算原数据各target类别的占比
target_ratio = df['target'].value_counts(normalize=True)
# 按占比分配每个类别的采样行数
sample_n_per_target = (target_ratio * total_sample_n).round().astype(int)

# 修正行数误差，避免分配后总条数和预期不一致
diff = total_sample_n - sample_n_per_target.sum()
if diff != 0:
    sample_n_per_target.iloc[0] += diff

# 分组执行抽样
sampled_df = df.groupby('target', group_keys=False).apply(
    lambda x: x.sample(n=sample_n_per_target[x.name], random_state=42)
)

方案说明 #

• 该方法逻辑简单高效，针对千万级行、数百列的大规模数据集也可以正常运行，无额外内存开销
• 每个target分组内的随机抽样天然保留了组内所有列的单变量分布、列间联合分布，不会出现分布偏移
• 可通过以下代码验证采样后的分布是否符合预期：

for col in sampled_df.columns:
    print(sampled_df.groupby(['target'])[col].value_counts(normalize=True))

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