按B列分层且按A列分组拆分Pandas DataFrame的方法

核心需求 #

• 同一A列取值的所有行只能全部进入训练集或测试集（按A分组拆分）
• 保持B列的类别分布比例（分层拆分）
• 允许删除部分行以调整比例符合要求

实现步骤 #

1. 统计分组与整体分布

先按A列分组，查看每个组内B列的类别分布，同时计算整体B列的分布比例，作为分层目标。

import pandas as pd
import numpy as np
import itertools

# 生成示例数据
size = 30
column_a_values = np.random.randint(1, 4, size=size)
column_b_values = np.random.choice(['x', 'y', 'z'], size=size)
column_c_values = np.random.rand(size)
data = {'A': column_a_values, 'B': column_b_values, 'C': column_c_values}
df = pd.DataFrame(data)

# 统计每个A组的B类别计数
group_b_dist = df.groupby('A')['B'].value_counts().unstack(fill_value=0)
group_b_dist['total'] = group_b_dist.sum(axis=1)
print("各A组的B类别分布：")
print(group_b_dist)

# 计算整体B列的分布比例（目标分层比例）
target_ratio = 0.7  # 训练集占比
overall_b_dist = df['B'].value_counts(normalize=True)
print("\n整体B列分布比例：")
print(overall_b_dist)

2. 选择符合分层要求的A组组合

遍历所有可能的A组组合，找到能让训练集B分布最接近整体分布、且占比接近目标比例的组合。如果A的取值较多，可改用贪心策略提升效率。

a_groups = df['A'].unique().tolist()
best_error = float('inf')
best_train_groups = []
tolerance = 0.05  # 比例允许的误差范围

# 遍历所有A组组合（适用于A取值较少的场景）
for k in range(1, len(a_groups)+1):
    for combo in itertools.combinations(a_groups, k):
        train_candidate = df[df['A'].isin(combo)]
        # 计算训练集B分布与整体分布的误差（平方差之和）
        train_b_dist = train_candidate['B'].value_counts(normalize=True).reindex(overall_b_dist.index, fill_value=0)
        error = ((train_b_dist - overall_b_dist)**2).sum()
        # 计算当前组合的训练集占比
        current_ratio = len(train_candidate) / len(df)
        
        # 筛选误差最小且比例符合要求的组合
        if abs(current_ratio - target_ratio) <= tolerance and error < best_error:
            best_error = error
            best_train_groups = list(combo)

# 如果未找到符合误差范围的组合，选择最接近目标比例的组合
if not best_train_groups:
    best_ratio_diff = float('inf')
    for k in range(1, len(a_groups)+1):
        for combo in itertools.combinations(a_groups, k):
            train_candidate = df[df['A'].isin(combo)]
            current_ratio = len(train_candidate) / len(df)
            ratio_diff = abs(current_ratio - target_ratio)
            if ratio_diff < best_ratio_diff:
                best_ratio_diff = ratio_diff
                best_train_groups = list(combo)

3. 拆分并调整数据集

根据选中的A组拆分训练集和测试集，若比例仍不符合要求，可在训练集/测试集内按B分层删除部分行。

# 初始拆分
train_df = df[df['A'].isin(best_train_groups)]
test_df = df[~df['A'].isin(best_train_groups)]

# 检查并调整比例（示例：训练集占比过高时删除部分行）
final_train_ratio = len(train_df) / len(df)
if final_train_ratio > target_ratio:
    # 计算需要删除的行数
    excess_rows = int(len(df) * (final_train_ratio - target_ratio))
    # 按B分层随机删除行，保持B分布比例
    train_df = train_df.groupby('B', group_keys=False).apply(
        lambda x: x.sample(n=len(x) - int(len(x)*excess_rows/len(train_df)), random_state=42)
    ).reset_index(drop=True)

# 输出结果
print("\n训练集B分布：")
print(train_df['B'].value_counts(normalize=True))
print("\n测试集B分布：")
print(test_df['B'].value_counts(normalize=True))
print(f"\n最终训练集占比：{len(train_df)/len(df):.2f}")

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