如何在Pandas DataFrame中按指定规则生成Assigned ID列

问题背景 #

现有包含员工合同信息的DataFrame，字段包括Name（员工姓名）、Employee Type（员工类型）、ID（合同ID）、Primary Contract?（是否为主要合同）。其中：

• 员工仅有1份合同时，该行标记为"Yes"
• 员工有多份合同时，仅1行标记为"Yes"，其余为"No"

示例输入数据：

Name    Employee Type   ID  Primary Contract?
Paulo   Employee        10  Yes
Paulo   Contractor      10  No
Paulo   Contractor      10  No
Paulo   Employee        10  No
Paulo   Employee        10  No
Pedro   Employee        11  No
Pedro   Contractor      11  Yes
Jonas   Contractor      16  Yes
Jonas   Contractor      16  No
Felipe  Non Worker      15  Yes
Maria   Employee        13  No
Maria   Contractor      13  Yes
Maria   Employee        13  No

生成规则 #

需要创建名为Assigned ID的列，规则如下：

• 标记为"Yes"的行：Assigned ID = 员工类型首字母 + ID编号
• 标记为"No"的行：
  1. 先按相同规则生成基础ID（员工类型首字母 + ID编号）
  2. 若该基础ID已存在，则添加"-2"，重复则依次递增后缀（如-3、-4等）
• 首字母不同的基础ID视为不同ID

预期输出 #

Name    Employee Type   ID  Primary Contract?   Assigned ID
Paulo   Employee        10  Yes                 E10
Paulo   Contractor      10  No                  C10
Paulo   Contractor      10  No                  C10-2
Paulo   Employee        10  No                  E10-2
Paulo   Employee        10  No                  E10-3
Pedro   Employee        11  No                  E11
Pedro   Contractor      11  Yes                 C11
Jonas   Contractor      16  Yes                 C16
Jonas   Contractor      16  No                  C16-2
Felipe  Non Worker      15  Yes                 N15
Maria   Employee        13  No                  E13
Maria   Contractor      13  Yes                 C13
Maria   Employee        13  No                  E13-2

解决方案 #

使用Pandas分组和序号分配实现，步骤如下：

1. 生成基础ID：提取员工类型首字母，拼接ID
2. 按基础ID分组，为每组内的行分配序号（保证主合同行排在组内首位）
3. 根据序号和是否为主要合同，生成最终的Assigned ID

实现代码：

import pandas as pd

# 构造示例数据
data = [
    ["Paulo", "Employee", 10, "Yes"],
    ["Paulo", "Contractor", 10, "No"],
    ["Paulo", "Contractor", 10, "No"],
    ["Paulo", "Employee", 10, "No"],
    ["Paulo", "Employee", 10, "No"],
    ["Pedro", "Employee", 11, "No"],
    ["Pedro", "Contractor", 11, "Yes"],
    ["Jonas", "Contractor", 16, "Yes"],
    ["Jonas", "Contractor", 16, "No"],
    ["Felipe", "Non Worker", 15, "Yes"],
    ["Maria", "Employee", 13, "No"],
    ["Maria", "Contractor", 13, "Yes"],
    ["Maria", "Employee", 13, "No"],
]
df = pd.DataFrame(data, columns=["Name", "Employee Type", "ID", "Primary Contract?"])

# 1. 生成基础ID：员工类型首字母 + ID
df["base_id"] = df["Employee Type"].str[0] + df["ID"].astype(str)

# 2. 按base_id分组，组内将主合同行排前面，生成序号
df = df.sort_values(by=["base_id", "Primary Contract?"], ascending=[True, False])
df["group_seq"] = df.groupby("base_id").cumcount() + 1

# 3. 生成最终的Assigned ID
def generate_assigned_id(row):
    if row["Primary Contract?"] == "Yes":
        return row["base_id"]
    else:
        return row["base_id"] if row["group_seq"] == 1 else f"{row['base_id']}-{row['group_seq']}"

df["Assigned ID"] = df.apply(generate_assigned_id, axis=1)

# 恢复原始行顺序（可选操作）
df = df.sort_index()

# 移除中间辅助列，输出结果
result = df.drop(columns=["base_id", "group_seq"])
print(result.to_string(index=False))

代码说明 #

• 基础ID生成：通过str[0]提取员工类型首字母，拼接转为字符串的合同ID
• 分组排序：按基础ID分组，同时将主合同行（Yes）排在组内最前面，确保主合同行对应组内序号1
• 最终ID生成：主合同行直接使用基础ID；非主合同行中，组内序号1的用基础ID，序号≥2的添加-序号后缀
• 顺序恢复：若需要保持和原始输入的行顺序一致，最后按索引排序即可

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