Pandas执行groupby分组时如何限制单个分组的最大大小

不需要手动遍历groupby返回的对象做二次拆分，直接通过构造复合分组键的方式就能一步完成需求，实现逻辑简洁，执行效率也更高。

核心思路：

• 先对目标分组列column_a做基础分组，用cumcount()给每个基础组内的行按顺序生成从0开始的连续序号
• 用组内序号除以阈值s做整数整除，得到的结果作为二级分组键：每连续s个同属一个column_a值的行会拿到相同的二级键，自然拆分出大小不超过s的子组
• 把原列和二级分组键组合作为最终分组依据即可

代码实现 #

import pandas as pd

# 示例数据
df = pd.DataFrame({'column_a': [1,1,1,2,3,3,3,3,3]})
s = 2 # 单组最大元素数阈值

# 构造二级分组键
sub_group_id = df.groupby('column_a').cumcount() // s

# 执行分组
for _, sub_df in df.groupby([df['column_a'], sub_group_id]):
    print(tuple(sub_df['column_a']))

运行结果 #

(1, 1)
(1,)
(2,)
(3, 3)
(3, 3)
(3,)

和预期分组结果完全一致。

补充说明 #

• 这种实现是Pandas原生的向量化操作，数据量较大时性能远高于手动循环拆分组的写法
• 分组后可以正常使用所有groupby支持的聚合、转换、过滤等接口，不需要做额外适配
• 如果需要拆分后的组不保留原同值的连续顺序，只需要在构造sub_group_id前先对DataFrame做对应规则的打乱即可

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