Pandas如何基于0/1标记列生成字符串标签列表列并构建目标数据框

完整实现代码 #

import pandas as pd

# 若你已经读取了数据到df变量，可直接跳过样例构造部分
df = pd.DataFrame([
    ['c24a1b14d253.jpg', 0, 1, 1],
    ['9ee905a41651.jpg', 0, 0, 1],
    ['3f58d128caf6.jpg', 0, 1, 0],
    ['4ce6599e7b20.jpg', 0, 0, 1],
    ['0def470360e4.jpg', 0, 0, 1]
], columns=['filename', 'diabetic retinopathy', 'glaucoma', 'others'])

# 筛选所有标签列（排除filename列）
label_cols = df.columns.drop('filename')

# 生成labels列：逐行提取值为1的列名转为列表
df['labels'] = df.apply(lambda row: label_cols[row == 1].tolist(), axis=1)

# 生成仅保留filename和labels的新数据框
result_df = df[['filename', 'labels']]

输出结果示例 #

运行后打印result_df可得到符合要求的结果：

filename               labels
0  c24a1b14d253.jpg  [glaucoma, others]
1  9ee905a41651.jpg            [others]
2  3f58d128caf6.jpg          [glaucoma]
3  4ce6599e7b20.jpg            [others]
4  0def470360e4.jpg            [others]

代码说明 #

• df.columns.drop('filename') 会自动排除文件名列，后续如果新增其他疾病标签列也不需要手动修改这段代码
• apply 方法中指定axis=1表示按行处理数据，对每一行用布尔索引筛选出值为1的对应列名，转为列表后赋值给labels列
• 最后直接提取指定两列即可得到目标数据框

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