如何以分层且按分区方式拆分Python Pandas DataFrame为N块？

核心思路 #

要同时满足分层分布近似一致、每个块的_PARTITION唯一值尽可能少的要求，采用分层内按分区聚合分配的策略：

1. 先以_STRATA为分组依据，保证每个分层的样本按比例分配到N个数据块；
2. 在每个分层内部，按_PARTITION聚合样本量，优先将大的分区整体分配到单个数据块，减少跨块的分区数量；
3. 处理样本量无法被N整除的边缘情况，将剩余样本均匀分散到各块，避免破坏分层比例。

实现代码 #

import pandas as pd
import numpy as np

def stratified_partition_split(df, n_splits, strata_col="_STRATA", partition_col="_PARTITION"):
    # 初始化每个拆分块的索引列表
    split_indices = [[] for _ in range(n_splits)]
    
    # 按分层列分组处理
    for strata, group in df.groupby(strata_col):
        # 按分区列聚合，计算每个分区的样本量（从大到小排序）
        partition_counts = group.groupby(partition_col).size().sort_values(ascending=False)
        
        # 计算当前分层的分配配额：基础量+余数
        total = len(group)
        base_quota = total // n_splits
        remainder = total % n_splits
        
        split_quota = [base_quota] * n_splits
        # 余数均匀分配给前remainder个块
        for i in range(remainder):
            split_quota[i] += 1
        
        # 按配额分配分区样本
        current_quota = split_quota.copy()
        for partition, cnt in partition_counts.items():
            # 优先分配给剩余配额最多的块
            target_idx = np.argmax(current_quota)
            if cnt <= current_quota[target_idx]:
                # 分区整体分配
                split_indices[target_idx].extend(group[group[partition_col] == partition].index.tolist())
                current_quota[target_idx] -= cnt
            else:
                # 分区过大时拆分分配
                part_indices = group[group[partition_col] == partition].index.tolist()
                # 先填满目标块剩余配额
                take = current_quota[target_idx]
                split_indices[target_idx].extend(part_indices[:take])
                part_indices = part_indices[take:]
                current_quota[target_idx] = 0
                
                # 剩余部分继续分配给配额充足的块
                while part_indices:
                    target_idx = np.argmax(current_quota)
                    take = min(current_quota[target_idx], len(part_indices))
                    split_indices[target_idx].extend(part_indices[:take])
                    part_indices = part_indices[take:]
                    current_quota[target_idx] -= take
    
    # 生成拆分后的DataFrame列表
    split_dfs = [df.loc[indices].copy() for indices in split_indices]
    return split_dfs

# 示例测试
if __name__ == "__main__":
    np.random.seed(42)
    test_data = {
        "_STRATA": np.random.choice(["X", "Y", "Z"], size=1200, p=[0.25, 0.45, 0.3]),
        "_PARTITION": np.random.choice([f"P{i}" for i in range(8)], size=1200),
        "feature": np.random.randn(1200)
    }
    df = pd.DataFrame(test_data)
    
    # 拆分为4个数据块
    result_splits = stratified_partition_split(df, n_splits=4)
    
    # 验证分层比例
    print("原数据集分层占比:")
    print(df["_STRATA"].value_counts(normalize=True).round(3))
    for idx, split_df in enumerate(result_splits):
        print(f"\n第{idx+1}块分层占比:")
        print(split_df["_STRATA"].value_counts(normalize=True).round(3))
    
    # 查看每个块的分区唯一值数量
    for idx, split_df in enumerate(result_splits):
        print(f"\n第{idx+1}块_PARTITION唯一值数量: {split_df['_PARTITION'].nunique()}")

关键说明 #

• 约束1&2严格满足：通过索引独立分配，确保所有数据块无重叠、并集完全覆盖原数据集；
• 约束3近似满足：每个分层按基础配额+余数的方式分配样本，分层比例与原数据集偏差极小；
• 约束4近似满足：优先将大分区整体分配到单个块，仅在分区样本量超过块剩余配额时才拆分，最大限度减少每个块的_PARTITION唯一值数量；
• 边缘情况处理：余数均匀分配解决样本量无法被N整除的问题，避免单块样本量偏差过大。

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