如何在Pandas中根据id是否含a和b生成指示列？

我有如下Pandas DataFrame：

import pandas as pd
df = pd.DataFrame({'id': [1,1,1,2,2,3,3,3],
                   'col': ['a','a','a','a','b','c','b','a']})

原始数据结构：

我希望新增一个indicator列，标记每个id对应的col中是否同时包含'a'和'b'，满足条件则标记为1，否则为0。期望输出如下：

result = pd.DataFrame({'id': [1,1,1,2,2,3,3,3],
                       'col': ['a','a','a','a','b','c','b','a'], 
                       'indicator': [0,0,0,1,1,1,1,1]})

对应的结果表格：

解决方案 #

方法1：使用groupby + transform #

通过分组后判断每个组是否同时包含'a'和'b'，再用transform将结果广播到每一行：

# 分组后检查每个组的col集合是否同时包含a和b
df['indicator'] = df.groupby('id')['col'].transform(
    lambda x: 1 if {'a', 'b'}.issubset(x.unique()) else 0
)

方法2：使用groupby + merge #

先对每个id计算是否满足条件，再合并回原DataFrame：

# 计算每个id的indicator值
id_indicator = df.groupby('id')['col'].apply(
    lambda x: 1 if {'a', 'b'}.issubset(x.unique()) else 0
).reset_index(name='indicator')

# 合并到原表
df = df.merge(id_indicator, on='id', how='left')

方法3：更高效的向量化实现 #

如果数据量较大，可通过crosstab生成标记矩阵，再判断条件：

# 生成每个id是否包含a和b的矩阵
cross = pd.crosstab(df['id'], df['col']).ge(1)
# 同时包含a和b则为1
cross['indicator'] = (cross['a'] & cross['b']).astype(int)
# 合并回原表
df = df.merge(cross['indicator'], on='id', how='left')

以上三种方法都能得到期望的结果，其中方法1代码最简洁，方法3在大数据集下效率更高。

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