如何拆分含日期与目标列的DataFrame，保留排序且各部分类别均衡

我有一个包含date和target（取值为0或1）两列、15万行的DataFrame，已按date排序。需要将其拆分为可调整数量的多个部分，满足以下要求：

• 每个部分中0和1的比例保持相等
• 保留初始排序
• 同一日期的行需归为同一部分，不能跨部分混合

解决思路 #

1. 先按date分组，统计每个日期下0和1的样本量，明确每个日期对全局正负样本的贡献占比
2. 按照预设的拆分数量，计算每个批次需要的正负样本目标量
3. 按日期顺序遍历，动态分配批次：每次尝试将当前日期加入当前批次，若加入后比例偏离目标过多，则切换到下一批次，确保同一日期的所有行都在同一个批次里
4. 最后给原DataFrame的每一行映射对应的批次编号

代码实现 #

import pandas as pd

def split_df_by_date_and_ratio(df, num_parts=2):
    # 按日期分组，统计每个日期的0和1样本数
    date_stats = df.groupby('date')['target'].agg(
        count_0=lambda x: (x == 0).sum(),
        count_1=lambda x: (x == 1).sum()
    ).reset_index()
    
    # 全局正负样本总数
    total_0 = date_stats['count_0'].sum()
    total_1 = date_stats['count_1'].sum()
    
    # 每个批次的目标正负样本量
    target_0_per_part = total_0 / num_parts
    target_1_per_part = total_1 / num_parts
    
    # 初始化累计计数、当前批次编号，以及日期-批次映射表
    cum_0 = 0
    cum_1 = 0
    current_part = 1
    date_to_part = {}
    
    # 遍历每个日期分配批次
    for _, date_row in date_stats.iterrows():
        # 模拟加入当前日期后的累计样本量
        temp_cum_0 = cum_0 + date_row['count_0']
        temp_cum_1 = cum_1 + date_row['count_1']
        
        # 如果加入后与当前批次的目标比例偏差过大，且还有剩余批次，就切换批次
        if (abs(temp_cum_0 / target_0_per_part - current_part) > 0.1 
            or abs(temp_cum_1 / target_1_per_part - current_part) > 0.1) and current_part < num_parts:
            current_part += 1
        
        date_to_part[date_row['date']] = current_part
        cum_0 += date_row['count_0']
        cum_1 += date_row['count_1']
    
    # 给原DataFrame添加批次列
    df['part'] = df['date'].map(date_to_part)
    
    # 打印每个批次的正样本比例，验证结果
    for part in range(1, num_parts+1):
        part_data = df[df['part'] == part]
        pos_ratio = part_data['target'].mean()
        print(f"批次{part}的正样本比例: {pos_ratio:.2%}")
    
    return df

# 示例调用（替换成你的DataFrame即可）
# df = pd.read_csv('your_data_file.csv')
# split_result = split_df_by_date_and_ratio(df, num_parts=3)

关键细节说明 #

• 日期完整性：通过先按日期分组的方式，从根源上避免了同一日期的行被拆分到不同批次的问题
• 比例控制：用相对误差判断是否切换批次，既保证了每个批次的比例接近目标，又能适配单个日期样本量较大的情况
• 排序保留：因为原DataFrame已经按date排序，分配批次时也是按日期顺序遍历，最终结果自然保留了初始排序

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