将DataFrame布尔列映射为分类列class的实现方法

给定如下DataFrame：

import pandas as pd

data = \
{'id': [1, 2, 3, 4, 5],
 'A': [1, 0, 0, 0, 0],
 'B': [0, 0, 1, 0, 0],
 'C': [0, 0, 0, 0, 1],
 'D': [0, 1, 0, 0, 0],
 'E': [0, 0, 0, 1, 0]}

df = pd.DataFrame(data)

需要新增一列class，规则为：

• 当A=1时，class取值0
• 当B=1时，class取值1
• 当C=1时，class取值2
• 当D=1时，class取值3
• 当E=1时，class取值4

解决方案 #

方法1：idxmax+映射字典 #

先提取每行中值为1的列名，再通过字典映射为目标数值：

class_map = {'A': 0, 'B': 1, 'C': 2, 'D': 3, 'E': 4}
# 仅针对A-E列操作，取每行最大值对应的列名后映射
df['class'] = df[['A', 'B', 'C', 'D', 'E']].idxmax(axis=1).map(class_map)

方法2：矩阵乘法dot #

直接用列对应的权重值与列值相乘求和，高效简洁：

# 对应A-E列的权重依次为0、1、2、3、4
weights = [0, 1, 2, 3, 4]
df['class'] = df[['A', 'B', 'C', 'D', 'E']].dot(weights)

预期输出 #

id  A  B  C  D  E  class
0   1  1  0  0  0  0      0
1   2  0  0  0  1  0      3
2   3  0  1  0  0  0      1
3   4  0  0  0  0  1      4
4   5  0  0  1  0  0      2

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