基于Pandas DataFrame按合同1所属公司创建'Primary Contract'列

需求说明 #

针对包含员工ID、合同编号、所属公司的DataFrame，为每个员工找到其合同编号为1对应的公司，创建Primary Contract列：当该行合同所属公司与该员工合同1的公司一致时，值为Yes，否则为No。

数据示例 #

原始DataFrame

ID  Contract Number Company
10000   1           Abc
10000   2           Zxc
10000   3           Abc
10001   1           Zxc
10002   2           Abc
10002   1           Cde
10002   3           Zxc

目标DataFrame

ID  Contract Number Company Primary Contract
10000   1            Abc           Yes
10000   2            Zxc           No
10000   3            Abc           Yes
10001   1            Zxc           Yes
10002   2            Abc           No
10002   1            Cde           Yes
10002   3            Zxc           No

最优实现方法 #

推荐使用向量化映射+布尔判断的方案，避免行级循环，效率更高，代码简洁：

1. 首先构建员工ID到其主合同公司（合同编号1）的映射关系
2. 通过映射将每个员工的主公司匹配到原表的每一行
3. 直接比较当前行公司与主公司，转换为Yes/No标记

代码实现

import pandas as pd

# 构建原始数据
df = pd.DataFrame({
    'ID': [10000, 10000, 10000, 10001, 10002, 10002, 10002],
    'Contract Number': [1, 2, 3, 1, 2, 1, 3],
    'Company': ['Abc', 'Zxc', 'Abc', 'Zxc', 'Abc', 'Cde', 'Zxc']
})

# 步骤1：提取每个ID对应的主合同公司（仅保留Contract Number=1的行）
primary_map = df[df['Contract Number'] == 1].set_index('ID')['Company']

# 步骤2：匹配主公司并生成标记列
df['Primary Contract'] = (df['Company'] == df['ID'].map(primary_map)).map({True: 'Yes', False: 'No'})

# 查看结果
print(df)

方案优势 #

• 完全利用Pandas向量化运算，比apply行级循环效率提升数倍，适合处理大规模数据
• 逻辑清晰，代码量少，易于维护

备选方案（GroupBy+Transform） #

如果需要更直观的分组逻辑，也可以使用GroupBy结合Transform实现：

def get_primary(group):
    # 取当前分组中Contract Number=1的公司（假设每个ID至少有1条合同1记录）
    return group[group['Contract Number'] == 1]['Company'].iloc[0]

# 生成每个行对应的主公司
df['Primary Company'] = df.groupby('ID').transform(get_primary)
# 生成标记列
df['Primary Contract'] = (df['Company'] == df['Primary Company']).map({True: 'Yes', False: 'No'})
# 清理临时列
df.drop('Primary Company', axis=1, inplace=True)

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