如何从DataFrame与标签列表生成分组二进制标签？

已知条件 #

• 标签列表：

labs = ['G1','G2','G3','G4','G5','G6','G7']

• 去重后的目标数据集：

group entity_label
0      0           G1
1      0           G2
3      1           G5
4      1           G1
5      2           G1
6      2           G2
7      2           G3

（可通过原df_test执行drop_duplicates(subset=['group','entity_label'], keep='first')得到）

需求 #

为每个group生成对应labs的二进制标签列表：列表中每个位置对应labs的标签，若该组包含此标签则标记为1，否则为0，最终得到如下结果：

group    entity_label_binary
0      0  [1, 1, 0, 0, 0, 0, 0]
1      1  [1, 0, 0, 0, 1, 0, 0]
2      2  [1, 1, 1, 0, 0, 0, 0]

实现方案 #

方案一：分组聚合+列表推导 #

通过分组聚合获取每个组的标签集合，再遍历labs生成二进制列表：

import pandas as pd

labs = ['G1','G2','G3','G4','G5','G6','G7']
df_test = pd.DataFrame({'group': [0,0,0,1,1,2,2,2,2],
                       'entity_label':['G1','G2','G2','G5','G1','G1','G2','G3','G3']})

# 1. 数据去重
df_unique = df_test.drop_duplicates(subset=['group','entity_label'], keep='first')

# 2. 分组聚合每个组的标签为集合
grouped_tags = df_unique.groupby('group')['entity_label'].agg(set).reset_index()

# 3. 生成二进制标签列表
grouped_tags['entity_label_binary'] = grouped_tags['entity_label'].apply(
    lambda tag_set: [1 if lab in tag_set else 0 for lab in labs]
)

# 4. 整理结果
result = grouped_tags.drop('entity_label', axis=1)
print(result)

方案二：交叉表法 #

利用pd.crosstab生成交叉统计表格，再转换为二进制列表，代码更简洁：

import pandas as pd

labs = ['G1','G2','G3','G4','G5','G6','G7']
df_test = pd.DataFrame({'group': [0,0,0,1,1,2,2,2,2],
                       'entity_label':['G1','G2','G2','G5','G1','G1','G2','G3','G3']})

# 1. 数据去重
df_unique = df_test.drop_duplicates(subset=['group','entity_label'], keep='first')

# 2. 生成交叉表，统计每组各标签的存在情况
cross_table = pd.crosstab(df_unique['group'], df_unique['entity_label'])
# 按labs的顺序调整列顺序，确保二进制列表对应正确
cross_table = cross_table[labs]

# 3. 将每行转换为二进制列表
cross_table['entity_label_binary'] = cross_table.apply(list, axis=1)

# 4. 整理结果
result = cross_table.reset_index()[['group', 'entity_label_binary']]
print(result)

两种方案运行后都会输出目标结果。

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