多标签图像分类任务中，如何为DataFrame生成含标签列表的新列？

问题场景 #

你有包含id和labels两列的DataFrame，需要新增all_labels列：当id存在重复时，将该id对应的所有标签以列表形式填入该列的首行，其余重复行的all_labels留空。

解决方案（基于Pandas） #

以下是可直接复用的代码步骤：

1. 导入依赖并创建示例DataFrame

import pandas as pd

# 模拟你的原始数据
df = pd.DataFrame({
    'id': ['x.jpg', 'x.jpg', 'y.jpg', 'z.jpg', 'z.jpg'],
    'labels': ['label_1', 'label_2', 'label_a', 'label_x', 'label_y']
})

2. 生成每个id对应的标签列表映射
   通过groupby聚合每个id下的所有标签为列表：

# 建立id到对应标签列表的映射
id_to_labels = df.groupby('id')['labels'].apply(list)

3. 合并映射到原DataFrame并处理空值
   先将映射关联到原表，再通过分组保留每组首行的标签列表，其余行设为空：

# 先给所有行添加对应的标签列表
df['all_labels'] = df['id'].map(id_to_labels)
# 仅保留每个id组的第一行的标签列表，其余行置空
df['all_labels'] = df.groupby('id')['all_labels'].transform(lambda x: x.iloc[0] if len(x) > 1 else None)

最终效果 #

处理后的DataFrame如下：

补充说明 #

• 如果你希望所有重复id的行都显示完整标签列表，只需跳过第三步，保留第二步的结果即可。
• 若你的DataFrame中id的重复顺序有特殊要求，可先对id和其他列排序后再执行上述步骤。

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