如何基于DataFrame现有列生成唯一子ID(sub ID)？

问题场景 #

现有一个包含Name、Primary row?、Employee Type、ID列的DataFrame，数据示例：

Name          Primary row?  Employee Type   ID
Paulo Cortez  Yes           Employee        100000
Paulo Cortez  No            Employee        100000
Joan San      Yes           Non-employee    100001
Felipe Castro Yes           Contractor      100002
Felipe Castro No            Employee        100002
Felipe Castro No            Contractor      100002

需要生成sub ID列，规则如下：

• 基础格式：提取Employee Type的首字母，拼接ID值
• 同一ID存在多条记录时：
  • Primary row?为Yes的行保持基础格式
  • Primary row?为No的行依次添加-2、-3这类后缀

预期结果：

Name          Primary row?  Employee Type   ID      sub ID
Paulo Cortez  Yes           Employee        100000  E100000
Paulo Cortez  No            Employee        100000  E100000-2
Joan San      Yes           Non-employee    100001  N100001
Felipe Castro Yes           Contractor      100002  C100002
Felipe Castro No            Employee        100002  E100002-2
Felipe Castro No            Contractor      100002  E100002-3

最优解法 #

使用pandas的分组和向量化操作实现，代码如下：

import pandas as pd

# 构造示例数据
data = {
    'Name': ['Paulo Cortez', 'Paulo Cortez', 'Joan San', 'Felipe Castro', 'Felipe Castro', 'Felipe Castro'],
    'Primary row?': ['Yes', 'No', 'Yes', 'Yes', 'No', 'No'],
    'Employee Type': ['Employee', 'Employee', 'Non-employee', 'Contractor', 'Employee', 'Contractor'],
    'ID': [100000, 100000, 100001, 100002, 100002, 100002]
}
df = pd.DataFrame(data)

# 1. 生成sub ID的基础前缀
df['prefix'] = df['Employee Type'].str[0] + df['ID'].astype(str)

# 2. 按ID分组，为每组内的行分配序号（从1开始）
df['seq'] = df.groupby('ID').cumcount() + 1

# 3. 根据Primary标记生成最终sub ID
df['sub ID'] = df.apply(
    lambda x: x['prefix'] if x['Primary row?'] == 'Yes' else f"{x['prefix']}-{x['seq']}",
    axis=1
)

# 清理中间列（可选）
df = df.drop(['prefix', 'seq'], axis=1)

print(df)

解法说明 #

1. 生成前缀：通过str[0]提取Employee Type首字母，和转成字符串的ID拼接，得到每个行的基础标识
2. 分组序号：用groupby('ID').cumcount() +1给同一ID下的行从1开始编号，Primary行对应序号1，后续非Primary行依次为2、3...
3. 生成最终sub ID：通过apply判断每行是否为Primary行，是则直接使用前缀，否则拼接前缀和序号

这种方法的优势在于：

• 利用pandas的向量化运算，性能远高于逐行循环，适合处理大规模数据
• 逻辑拆解清晰，代码可读性强，后续需求调整时容易修改

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