如何让Pandas的groupby().cumcount()忽略DataFrame中的非重复项？

需求说明 #

需要对DataFrame中重复的NameID添加序号标记（如WS3_1、WS3_2），而非重复的NameID保持原样，同时保留所有记录。

输入DataFrame #

原代码问题 #

你提供的原代码会给所有NameID添加序号后再尝试替换_0，但cumcount().add(1)会让非重复项的序号为1，导致最终生成WS1_1这类不符合预期的结果。

df = pd.DataFrame({
    'NameID': ['WS1', 'WS2', 'WS3', 'WS3', '01-1000', '01-1000', '01-1000', '01-1001'],
    'Type': ['A', 'B', 'C1', 'C2', 'D1', 'D1', 'D1', 'E'],
    'Value':['1', '2', '3', '4', '5', '5', '5', '6']
})   

df['NameID'] = df['NameID'].str.cat(df.groupby('NameID').cumcount().add(1).astype(str), sep='_').str.replace("_0", "")

解决方案 #

通过先判断NameID的出现次数，仅对重复项添加序号标记，高效实现需求：

import pandas as pd

df = pd.DataFrame({
    'NameID': ['WS1', 'WS2', 'WS3', 'WS3', '01-1000', '01-1000', '01-1000', '01-1001'],
    'Type': ['A', 'B', 'C1', 'C2', 'D1', 'D1', 'D1', 'E'],
    'Value':['1', '2', '3', '4', '5', '5', '5', '6']
})

# 生成每个组的序号（从1开始）
cum_counts = df.groupby('NameID').cumcount().add(1).astype(str)
# 判断哪些NameID是重复的（出现次数>1）
is_duplicated = df.groupby('NameID')['NameID'].transform('size') > 1
# 仅对重复的NameID拼接序号，非重复项保持原样
df['NameID'] = df['NameID'].mask(is_duplicated, df['NameID'] + '_' + cum_counts)

print(df)

执行后的结果DataFrame #

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