在Python Pandas中按行布尔值聚合对应列名的实现方案

需求说明 #

扫描Pandas DataFrame，识别每行中值为True的所有列，将这些列名聚合到一个新的输出列中。

示例输入 #

a       b       c
True    False   True
True    False   False

期望输出 #

a       b       c     output
True    False   True  ['a', 'c']
True    False   False ['a']

现有尝试（扩展性差） #

之前用np.where实现，但每新增一列就要添加对应代码，扩展性极差：

import numpy as np
import pandas as pd

df['output_1'] = np.where(df['a']==True, 'a', '')
df['output_2'] = np.where(df['b']==True, 'b', '')
df['output_3'] = np.where(df['c']==True, 'c', '')
df['output'] = df['output_1'] + df['output_2'] + df['output_3']

更优解决方案 #

方法1：apply + 列表推导（推荐） #

通过axis=1逐行遍历，筛选出值为True的列名，无需硬编码列名，新增列时无需修改代码：

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame({
    'a': [True, True],
    'b': [False, False],
    'c': [True, False]
})

df['output'] = df.apply(lambda row: [col for col, val in row.items() if val], axis=1)

方法2：dot方法（适合字符串拼接场景） #

如果需要先得到逗号分隔的字符串再转列表，可以用这种方式，不过性能略低于方法1：

df['output'] = df.dot(df.columns + ',').str.rstrip(',').str.split(',')

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