Pandas如何按指定列分组且每组固定n行 可从原DataFrame删除分组数据

你持有一个规模为数千行的pandas DataFrame，样例数据结构如下：

a   b   c
x   2   3
y   2   3
z   3   2
w   1   5
# 共数千条记录

最初你直接使用b、c列作为分组键遍历：

for x,y in df.groupby(['b','c']):
    print(y)

该方式会将相同b、c值的所有记录归为同一组，无法满足拆分要求。以n=1（单组最多1行）为例，默认分组结果如下：

a   b   c
x   2   3
y   2   3

a   b   c
z   3   2

a   b   c
w   1   5

你期望的效果是同一b、c组合下超过n行的部分自动拆分为新分组，n=1时期望输出为：

a   b   c
x   2   3

a   b   c
y   2   3

a   b   c
z   3   2

a   b   c
w   1   5

同时需要支持将已提取的分组数据从原始DataFrame中删除。

核心思路是在b、c基础分组之上，给组内每条记录计算子分组编号：使用组内累计序号对n做整除取整，将子分组编号与b、c共同作为分组键，即可自动完成超量记录的拆分，全程使用pandas内置向量化方法，性能远高于逐行遍历。

完整拆分代码 #

import pandas as pd

# 1. 初始化示例数据（替换为你自己的DataFrame即可）
df = pd.DataFrame({
    'a': ['x', 'y', 'z', 'w'],
    'b': [2, 2, 3, 1],
    'c': [3, 3, 2, 5]
})

n = 1  # 自定义每个分组的最大行数
# 2. 给每个b、c组内的记录生成子分组ID
df['sub_group_id'] = df.groupby(['b', 'c']).cumcount() // n
# 3. 按复合键分组，得到所有符合行数要求的分组
groups = [g.drop(columns=['sub_group_id']) for _, g in df.groupby(['b', 'c', 'sub_group_id'])]

# 遍历验证分组结果
for g in groups:
    print(g, '\n')

运行后输出与你期望的n=1拆分结果完全一致。调整n的数值即可灵活控制单组最大行数，例如n=2时，同一b、c下每2行会被归为一个独立分组。

提取分组后删除原表对应数据 #

不要在循环中逐行执行drop操作（性能极低），直接通过布尔索引筛选即可快速完成提取和原表裁剪：

# 示例：提取每个b、c组合下的第一组（前n行数据）
extracted_data = df[df['sub_group_id'] == 0].drop(columns=['sub_group_id'])
# 原表仅保留未被提取的剩余数据
df = df[df['sub_group_id'] != 0].drop(columns=['sub_group_id'])

如果需要按顺序逐批提取分组，每处理完一批分组，只需要通过行索引匹配筛选剩余数据即可，全流程无冗余遍历，处理十万级以上数据也能保持很高的运行效率。

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