Python中打乱后按组权重分层拆分50万条分组数据集的方法

针对你的需求，可以用pandas + numpy组合实现，既能保证数据打乱，又能按分组占比拆分成指定数量的数据集，每份都包含所有分组。以下是具体步骤和代码示例：

核心思路 #

1. 先整体打乱数据集，保证随机性；
2. 按分组逐个处理，对每个组按目标权重（或等比例）拆分；
3. 将每个组的拆分片段对应合并，得到最终的多个小型数据集。

代码实现 #

假设你的数据集是pandas.DataFrame格式，包含一个用于分组的列（比如group_id）：

import pandas as pd
import numpy as np

# 1. 加载并打乱数据集（替换成你的数据加载逻辑）
df = pd.read_csv("your_data.csv")
# 整体打乱，random_state可选，用于复现结果
df_shuffled = df.sample(frac=1, random_state=42).reset_index(drop=True)

# 2. 定义拆分参数
split_count = 10  # 要拆成10份
# 可选：自定义权重（总和需为1，示例为等比例拆分）
weights = [1/split_count] * split_count

# 3. 初始化存储拆分结果的字典
split_datasets = {f"dataset_{i+1}": [] for i in range(split_count)}

# 4. 按分组拆分并分配
for group_name, group_data in df_shuffled.groupby("group_id"):
    group_size = len(group_data)
    
    # 计算当前组在每个拆分数据集中的行数
    split_sizes = [round(group_size * w) for w in weights]
    # 调整拆分大小，确保总和等于当前组的行数
    size_diff = group_size - sum(split_sizes)
    if size_diff > 0:
        # 余数依次加到前size_diff个数据集
        for i in range(size_diff):
            split_sizes[i] += 1
    elif size_diff < 0:
        # 缺额依次从后size_diff个数据集减去
        for i in range(-size_diff):
            split_sizes[-(i+1)] -= 1
    
    # 拆分当前组
    group_splits = np.array_split(group_data, np.cumsum(split_sizes)[:-1])
    
    # 将拆分片段分配到对应数据集
    for idx, split in enumerate(group_splits):
        split_datasets[f"dataset_{idx+1}"].append(split)

# 5. 合并每个数据集的片段
for name in split_datasets:
    split_datasets[name] = pd.concat(split_datasets[name], ignore_index=True)

方法优势 #

• 支持任意数量的拆分（不限于10份），且每份都包含所有分组；
• 可自定义拆分权重（比如某份占20%，其余各占8%），也支持等比例拆分；
• 处理50万条数据效率较高：pandas和numpy的操作均为向量化，避免低效循环；
• 先整体打乱数据，保证拆分后的数据集随机性。

对比你尝试过的方法 #

• array_split：仅做无差别拆分，不考虑分组，无法保证每份包含所有分组；
• scikit-learn的分层拆分工具（如StratifiedShuffleSplit）：仅支持训练/测试集的二分拆分，无法直接生成多份数据集，且多次拆分可能导致数据重叠。

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