如何用Pandas将多列合并为以列名为标签的category列？

原始数据 #

ID      A      B      C
0   0   True  False  False
1   1  False   True  False
2   2  False  False   True

目标结果 #

ID group
0   0     A
1   1     B
2   2     C

需求说明 #

• 用原布尔列的列名作为group列的分类标签
• 每行最多仅存在一个True值

最小可复现代码 #

#!/usr/bin/env python3
import pandas as pd

df = pd.DataFrame({
    'ID': range(3),
    'A': [True, False, False],
    'B': [False, True, False],
    'C': [False, False, True]
})

result = pd.DataFrame({
    'ID': range(3),
    'group': ['A', 'B', 'C']
})
result.group = result.group.astype('category')

print(df)
print(result)

优雅的Pandas原生实现方法 #

方法1：idxmax（最简洁高效） #

利用布尔值等价于1/0的特性，idxmax会直接返回每行最大值对应的列名：

# 提取布尔列，逐行找到True对应的列名，添加为group列并转为分类类型
result = df.assign(
    group=df[['A', 'B', 'C']].idxmax(axis=1)
).astype({'group': 'category'})[['ID', 'group']]

方法2：melt+筛选（通用场景友好） #

将宽表转为窄表后筛选出布尔值为True的行，适合后续需要扩展的场景：

# 宽表转窄表，保留ID，将列名转为group列，布尔值转为flag列
melted_df = df.melt(id_vars='ID', var_name='group', value_name='flag')
# 筛选flag为True的行，去掉多余列并转换类型
result = melted_df[melted_df['flag']].drop('flag', axis=1).astype({'group': 'category'})

方法3：矩阵点积dot（灵活的数值实现） #

通过布尔值与列名的点积直接得到目标标签：

# 布尔列与列名做点积，得到每行对应的标签
result = df.assign(
    group=df[['A', 'B', 'C']].dot(df[['A', 'B', 'C']].columns)
).astype({'group': 'category'})[['ID', 'group']]

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