如何在Pandas DataFrame中按ID生成最多5行的批次分组ID

要实现每个id对应的行按最多5行一组划分，并生成全局连续的grp_id，可以按以下步骤操作：

1. 先为每个id内部的行生成组内批次号：

   • 用groupby('id').cumcount()获取每个行在对应id组内的索引（从0开始）
   • 对索引做整数除法//5，再加1，得到每个id内部的批次号（每5行一个批次）

2. 基于id和内部批次号的组合，生成全局连续的grp_id：

   • 将id和内部批次号作为联合分组键，用groupby([...]).ngroup()生成全局组索引，再加1得到最终的grp_id

完整代码 #

import pandas as pd

# 原始DataFrame
df = pd.DataFrame([[1, 1],
                    [2, 1],
                    [3, 1],
                    [4, 1],
                    [5, 1],
                    [6, 1],
                    [7, 1],
                    [8, 2],
                    [9, 2],
                    [10, 3],
                    [11, 3],
                   [12, 3],
                   [13, 4],
                   [14, 5],
                   [15, 5],
                   [16, 5],
                   [17, 5],
                   [18, 5],
                   [19, 5],
                   [20, 5]])
df.columns = ['ln_num', 'id']

# 步骤1：生成每个id内部的批次号
df['inner_grp'] = df.groupby('id').cumcount() // 5 + 1

# 步骤2：生成全局连续的grp_id
df['grp_id'] = df.groupby(['id', 'inner_grp']).ngroup() + 1

# 可选：删除中间列inner_grp
df = df.drop('inner_grp', axis=1)

print(df)

结果验证 #

运行后生成的grp_id会和你给出的预期输出一致：

• id=1的7行分为2个内部批次，对应grp_id=1和2
• id=2的2行对应grp_id=3
• id=3的5行对应grp_id=4
• id=1后续的3行（ln_num=12-14）属于id=1的第2个内部批次，对应grp_id=2
• id=1最后1行（ln_num=15）属于id=1的第3个内部批次，对应grp_id=5
• 以此类推，所有组的grp_id全局连续递增

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