如何为DataFrame中按Name分组的数据生成连续批次编号

问题场景 #

原DataFrame结构：

Name  Age
0    U   20
1    U   20
2    U   20
3    U   18
4    I   45
5    I   68
6    I    8
7    D    7
8    D    6
9    I   89

需求 #

• 设置批次大小为3，新增Batch列，批次号从1开始全局连续递增
• 同一Name值内，每满3条数据则批次号加1
• 期望输出：

Name  Age  Batch
0    U   20   1
1    U   20   1
2    U   20   1
3    U   18   2
4    I   45   3
5    I   68   3
6    I    8   3
7    D    7   4
8    D    6   4
9    I   89   5

原代码resu['B'] = np.divmod(np.arange(len(resu)),3)[0]+1未考虑Name分组，无法满足需求。

解决方案代码 #

import pandas as pd
import numpy as np

# 构造原数据（已有DataFrame可跳过此步）
df = pd.DataFrame({
    'Name': ['U', 'U', 'U', 'U', 'I', 'I', 'I', 'D', 'D', 'I'],
    'Age': [20, 20, 20, 18, 45, 68, 8, 7, 6, 89]
})

# 1. 生成每个Name分组内的行索引（从0开始）
df['group_idx'] = df.groupby('Name').cumcount()

# 2. 计算每个组内的批次编号（从0开始）
df['group_batch'] = df['group_idx'] // 3

# 3. 将(Name, group_batch)组合映射为全局连续批次号（从1开始）
df['Batch'] = df.groupby(['Name', 'group_batch']).ngroup() + 1

# 移除中间辅助列
df = df.drop(['group_idx', 'group_batch'], axis=1)

print(df)

代码逻辑说明 #

1. 组内行索引：groupby('Name').cumcount()为每个Name分组内的每行分配从0开始的连续序号，用于定位行在组内的位置。
2. 组内批次划分：group_idx // 3将组内每3行划分为一个批次，得到组内的批次编号（从0开始）。
3. 全局批次编号：groupby(['Name', 'group_batch']).ngroup()将每个Name与组内批次的组合视为唯一单元，生成全局连续的编号，加1后让批次号从1开始。
4. 清理中间列：移除辅助计算的列，得到最终结果。

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