如何用Pandas按场景实现type与month均衡的条件抽样？

针对按Scenario抽样，要求Month、Type维度尽可能均衡的需求，以下是基于Pandas的具体实现方案：

核心逻辑 #

1. 按Scenario单独处理每个分组的抽样配额
2. 分两种场景分配抽样数量：
   • 当抽样总数≥当前Scenario的Month唯一值数量：优先保证每个Month下的Type样本数均衡，剩余样本补到样本量较多的子组
   • 当抽样总数<当前Scenario的Month唯一值数量：仅保证Type维度的样本均衡，剩余样本补到样本量较多的Type组
3. 所有分组抽样完成后合并结果

代码实现 #

import pandas as pd

def balanced_sample(df_business, df_sample_quota):
    # 初始化结果DataFrame
    result = pd.DataFrame(columns=df_business.columns)
    
    # 遍历每个Scenario的抽样配额
    for scenario, quota in df_sample_quota.set_index('scenario')['sample_num'].items():
        # 筛选当前Scenario的业务数据
        subset = df_business[df_business['scenario'] == scenario].copy()
        if subset.empty:
            continue
        
        month_unique_count = subset['month'].nunique()
        type_unique_count = subset['type'].nunique()
        
        if quota >= month_unique_count:
            # 场景1：样本量足够覆盖所有Month，保证Month+Type均衡
            # 统计每个Month+Type组合的可用样本数
            group_stats = subset.groupby(['month', 'type']).size().reset_index(name='available')
            
            # 先按Month均分基础配额
            per_month_quota = quota // month_unique_count
            remaining_after_month_split = quota % month_unique_count
            
            # 每个Month内按Type均分配额
            group_stats['quota'] = per_month_quota // type_unique_count
            remaining_per_month = per_month_quota % type_unique_count
            
            # 处理每个Month内的剩余配额，优先分给样本多的Type
            for month in group_stats['month'].unique():
                month_group = group_stats[group_stats['month'] == month]
                top_indices = month_group.sort_values('available', ascending=False).head(remaining_per_month).index
                group_stats.loc[top_indices, 'quota'] += 1
            
            # 处理全局剩余配额（Month均分后的剩余）
            if remaining_after_month_split > 0:
                sorted_groups = group_stats.sort_values(['quota', 'available'], ascending=[True, False])
                global_top_indices = sorted_groups.head(remaining_after_month_split).index
                group_stats.loc[global_top_indices, 'quota'] += 1
            
            # 确保配额不超过实际可用样本数
            group_stats['quota'] = group_stats.apply(lambda x: min(x['quota'], x['available']), axis=1)
            
            # 执行抽样
            sampled_parts = []
            for _, row in group_stats.iterrows():
                part = subset[(subset['month'] == row['month']) & (subset['type'] == row['type'])].sample(
                    n=row['quota'], random_state=42
                )
                sampled_parts.append(part)
            sampled_subset = pd.concat(sampled_parts)
        
        else:
            # 场景2：样本量不足覆盖所有Month，仅保证Type均衡
            type_stats = subset.groupby('type').size().reset_index(name='available')
            
            # 按Type均分基础配额
            per_type_quota = quota // type_unique_count
            remaining_after_type_split = quota % type_unique_count
            
            type_stats['quota'] = per_type_quota
            # 剩余配额分给样本多的Type
            top_type_indices = type_stats.sort_values('available', ascending=False).head(remaining_after_type_split).index
            type_stats.loc[top_type_indices, 'quota'] += 1
            
            # 确保配额不超过实际可用样本数
            type_stats['quota'] = type_stats.apply(lambda x: min(x['quota'], x['available']), axis=1)
            
            # 执行抽样
            sampled_parts = []
            for _, row in type_stats.iterrows():
                part = subset[subset['type'] == row['type']].sample(n=row['quota'], random_state=42)
                sampled_parts.append(part)
            sampled_subset = pd.concat(sampled_parts)
        
        # 合并当前Scenario的抽样结果
        result = pd.concat([result, sampled_subset], ignore_index=True)
    
    return result

关键说明 #

• 用random_state=42固定随机种子，确保抽样结果可复现
• 每次分配配额后都会做min(quota, available)校验，避免抽取超过实际存在的样本
• 剩余配额优先分配给样本量较多的子组，最大化均衡性
• 空Scenario会自动跳过，避免报错

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