根据另一个列值 pandas 数据帧创建"Yes"列

Question

假设我有一个包含员工id、合同号和他们工作的公司的数据框架。每个员工可以为同一家公司甚至不同的公司签订尽可能多的合同:

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

我需要找到一种方法来识别合同号的公司&;1&;每个ID，然后创建一个列"Primary Contract"它将被设置为&;yes &;如果该合同与合同号1的公司属于同一家公司，则产生此数据框:

ID  Contract Number Company Primary Compay
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

实现它的最好方法是什么?

Answer 1

您可以将groupby.apply与isin和numpy.where一起使用:

df['Primary Company'] = np.where(
df.groupby('ID', group_keys=False)
.apply(lambda g: g['Company'].isin(g.loc[g['Contract Number'].eq(1), 'Company'])
),
'Yes', 'No'
)

输出:

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

如果您可以使用布尔值(True/False)而不是'Yes'/'No':

df['Primary Company'] = (
df.groupby('ID', group_keys=False)
.apply(lambda g: g['Company'].isin(g.loc[g['Contract Number'].eq(1), 'Company']))
)

Answer 2

筛选Contract Number为1的行，在DataFrame.merge中使用左连接，比较indicator=True参数生成的_merge列:

mask = (df.merge(df[df['Contract Number'].eq(1)],
how='left', on=['ID','Company'], indicator=True)['_merge'].eq('both'))
df['Primary Company'] = np.where(mask, 'Yes','No')
print (df)
ID  Contract Number Company Primary Company
0  10000                1     Abc             Yes
1  10000                2     Zxc              No
2  10000                3     Abc             Yes
3  10001                1     Zxc             Yes
4  10002                2     Abc              No
5  10002                1     Cde             Yes
6  10002                3     Zxc              No

另一个想法是比较MultiIndex与Index.isin:

idx = df[df['Contract Number'].eq(1)].set_index(['ID','Company']).index
df['Primary Company'] = np.where(df.set_index(['ID','Company']).index.isin(idx),
'Yes','No')
print (df)
ID  Contract Number Company Primary Company
0  10000                1     Abc             Yes
1  10000                2     Zxc              No
2  10000                3     Abc             Yes
3  10001                1     Zxc             Yes
4  10002                2     Abc              No
5  10002                1     Cde             Yes
6  10002                3     Zxc              No