如何高效自动拆分Pandas DataFrame为无重复ID的n个分块？

np.array_split直接按行均分的方式会破坏同ID的完整性，要满足「同ID不跨块」的要求，核心是先对唯一ID进行分组，再映射回原DataFrame提取对应行。

基础实现方案 #

先提取所有不重复的ID，将这些ID拆分为4组，再根据每组ID筛选原DataFrame的行：

import pandas as pd
import numpy as np

# 构造示例DataFrame
data = {'ID': [1,1,1,20,20,30,50,50,51,60,60,300,300,302,302,500]}
df = pd.DataFrame(data)

# 提取唯一ID并拆分为4组
unique_ids = df['ID'].unique()
id_groups = np.array_split(unique_ids, 4)

# 生成拆分后的DataFrame列表
df_splits = [df[df['ID'].isin(group)] for group in id_groups]

# 验证结果（可选）
for idx, split_df in enumerate(df_splits, 1):
    print(f"分块{idx}包含的ID：{split_df['ID'].unique()}")
    print(f"分块{idx}的行数：{len(split_df)}\n")

均衡优化方案 #

如果希望各分块的行数尽量接近（避免某块过大/过小），可以先按ID的出现次数排序，再轮流分配ID到不同组：

# 按ID出现次数从多到少排序
id_counts = df['ID'].value_counts().sort_values(ascending=False)
sorted_ids = id_counts.index.tolist()

# 轮流分配ID到4个组，平衡各块行数
id_groups_balanced = [[], [], [], []]
for idx, id_val in enumerate(sorted_ids):
    id_groups_balanced[idx % 4].append(id_val)

# 生成均衡拆分后的DataFrame列表
df_splits_balanced = [df[df['ID'].isin(group)] for group in id_groups_balanced]

# 验证均衡结果（可选）
for idx, split_df in enumerate(df_splits_balanced, 1):
    print(f"均衡分块{idx}包含的ID：{split_df['ID'].unique()}")
    print(f"均衡分块{idx}的行数：{len(split_df)}\n")

两种方案都能保证同ID的所有行只出现在一个分块中，基础方案实现简单，均衡方案能让各分块的行数更接近。

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