行数无法被组数整除时，如何在DataFrame中按周期生成分组列？

我来帮你搞定这个问题——要给每个周期内的公司分组，当行数没法被组数整除时，让前两组各分2行、最后一组分3行（就像你例子里7家公司的情况），用Pandas就能轻松实现，下面是完整方案：

步骤1：构造示例数据 #

先模拟你说的场景：7家公司、2个周期的DataFrame：

import pandas as pd
import numpy as np

# 构造示例数据
data = {
    'company': ['A', 'B', 'C', 'D', 'E', 'F', 'G'] * 2,
    'period': ['P1'] * 7 + ['P2'] * 7
}
df = pd.DataFrame(data)

步骤2：定义分组函数 #

写一个函数，专门处理单个周期内的分组逻辑，核心是计算每组应分配的行数，处理整除不了的情况：

def assign_groups(group, num_groups=3):
    # 获取当前周期的公司数量
    total_rows = len(group)
    # 计算每组基础行数，以及剩余未分配的行数
    base_size, remainder = divmod(total_rows, num_groups)
    
    # 初始化每组的行数：所有组先分配基础行数
    group_sizes = [base_size] * num_groups
    # 把剩余的行数分配给最后几个组（这里余数是1，所以最后1组多1行）
    for i in range(remainder):
        group_sizes[-(i+1)] += 1
    
    # 生成对应的分组标签，用numpy重复数组快速生成
    group['res'] = np.repeat(range(1, num_groups+1), group_sizes)
    return group

步骤3：应用分组逻辑到每个周期 #

用groupby按周期分组，然后把上面的函数应用到每个分组：

# 按周期分组处理，然后重置索引
result_df = df.groupby('period').apply(assign_groups).reset_index(drop=True)

验证结果 #

打印result_df就能看到预期的res列了：

company period  res
0        A     P1    1
1        B     P1    1
2        C     P1    2
3        D     P1    2
4        E     P1    3
5        F     P1    3
6        G     P1    3
7        A     P2    1
8        B     P2    1
9        C     P2    2
10       D     P2    2
11       E     P2    3
12       F     P2    3
13       G     P2    3

完全符合你要求的前两组各2行、最后一组3行的规则~

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