如何在Pandas DataFrame中按主工作位置条件新增‘Primary Location?’列？

我有一个包含员工ID、所从事工作、工作地点以及标记该工作是否为员工主工作的DataFrame：

ID  Job     Location        Main Job?
10  Driver  Peru street     Yes
11  Nurse   Jorge street    No
11  Helper  Jorge street    Yes
11  Driver  Phoenix av.     No
11  Admin   Jorge street    No
12  Driver  Yard av.        Yes
12  Helper  Yard av.        No
12  Nurse   Jorge street    No
13  Driver  Peru street     Yes
13  Helper  Jorge street    No

已知每个唯一ID对应的“Main Job?”列仅有一项标记为“Yes”，现在需要新增一列Primary Location?：当员工的工作地点与其主工作的地点相同时标记为“Yes”，否则标记为“No”，最终得到如下结果：

ID  Job     Location       Main Job?    Primary Location?
10  Driver  Peru street     Yes         Yes
11  Nurse   Jorge street    No          Yes
11  Helper  Jorge street    Yes         Yes
11  Driver  Phoenix av.     No          No
11  Admin   Jorge street    No          Yes
12  Driver  Yard av.        Yes         Yes
12  Helper  Yard av.        No          Yes
12  Nurse   Jorge street    No          No
13  Driver  Peru street     Yes         Yes
13  Helper  Jorge street    No          No

用Pandas可以高效完成这个需求，核心思路是先提取每个员工ID对应的主工作地点，再将该地点与员工所有工作记录的地点做比对：

方法一：向量化操作（大数据集首选） #

这种方法利用Pandas底层优化，避免逐行遍历，性能更优：

# 提取每个ID对应的主工作地点
main_loc_map = df[df['Main Job?'] == 'Yes'].set_index('ID')['Location']
# 将主地点映射到原DataFrame
df = df.merge(main_loc_map.rename('Main_Location'), on='ID', how='left')
# 对比生成目标列
df['Primary Location?'] = (df['Location'] == df['Main_Location']).map({True: 'Yes', False: 'No'})
# 可选：删除临时辅助列
df = df.drop('Main_Location', axis=1)

方法二：逐行判断（小数据集可用） #

如果数据集不大，也可以用apply逐行处理，逻辑更直观：

# 先构建ID到主地点的字典映射
main_loc_dict = df[df['Main Job?'] == 'Yes'].set_index('ID')['Location'].to_dict()
# 逐行判断生成新列
df['Primary Location?'] = df.apply(lambda x: 'Yes' if x['Location'] == main_loc_dict[x['ID']] else 'No', axis=1)

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