如何为Pandas分组结果集添加自增GROUP_ID列？

修正后的代码 #

import pandas as pd
import numpy as np

# 读取数据
df = pd.read_csv('dups_check_group_v1.csv', encoding='unicode_escape', usecols=['ID','ENTITY_NAME','ENTITY_VALUE','SECTION_GROUP','DOC_ID'])

# 分离SECTION_GROUP为空的行
mask = df['SECTION_GROUP'].isna()
rest = df[mask]

# 按DOC_ID和SECTION_GROUP联合分组，确保每个文档下的每个组独立
grouped = df[~mask].groupby(['DOC_ID', 'SECTION_GROUP'])

frames = []
group_id_counter = 1

# 遍历每个分组，为当前组和对应文档的空行分配相同GROUP_ID
for (doc_id, _), group in grouped:
    # 为当前组添加GROUP_ID
    group['GROUP_ID'] = group_id_counter
    # 筛选当前文档对应的空行并添加GROUP_ID
    current_rest = rest[rest['DOC_ID'] == doc_id].copy()
    current_rest['GROUP_ID'] = group_id_counter
    # 将组和空行加入列表
    frames.append(group)
    frames.append(current_rest)
    # 递增GROUP_ID
    group_id_counter += 1

# 合并所有数据并按DOC_ID排序
out = pd.concat(frames).sort_values('DOC_ID')
print(out)

代码说明 #

1. 联合分组：改用['DOC_ID', 'SECTION_GROUP']分组，保证不同文档下的同名组（如GROUP 1）被视为独立块，避免跨文档的组混淆。
2. GROUP_ID分配：遍历每个分组时，为当前组的所有行和对应文档的空行分配同一个GROUP_ID，确保同一逻辑块的行共享ID。
3. 避免重复空行：仅筛选当前文档对应的空行进行拼接，解决原代码中空行重复出现的问题。

输出示例 #

执行后将得到如下结构的结果（符合期望的GROUP_ID）：

ID ENTITY_NAME   ENTITY_VALUE SECTION_GROUP  DOC_ID  GROUP_ID
0   1     dNumber  U220059090(C)       GROUP 1      40         1
1   2       tDate       6-Dec-22       GROUP 1      40         1
4   5    sCompany             bp           NaN      40         1
2   3     dNumber  U220059090(C)       GROUP 2      40         2
3   4       tDate       6-Dec-22       GROUP 2      40         2
4   5    sCompany             bp           NaN      40         2
5   6     dNumber  U220059090(C)       GROUP 1      42         3
6   7       tDate       6-Dec-22       GROUP 1      42         3
9  10    sCompany             bp           NaN      42         3
7   8     dNumber  U220059090(C)       GROUP 2      42         4
8   9       tDate       6-Dec-22       GROUP 2      42         4
9  10    sCompany             bp           NaN      42         4

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