将独热标签转分类标签：Pandas报KeyError: pos问题求助

我有一个CSV文件，包含对应ID的4类独热标签，示例结构如下：

ids   A  B  C D
1     0  1  0 0
2     0  0  1 0
3     1  0  0 0
...
10000

我希望新增一列pos来存储对应的分类，示例效果如下：

ids pos
 1   B
 2   C
 3   A

我使用Pandas进行处理，但出现了KeyError: pos错误。附上我的代码：

import pandas as pd

df=pd.read_csv("ABC.csv")

cols=['A','B','C','D']
df['arr']=df[cols].values.tolist()
print(df.head())

for ind in df.head().index:
    print(df['arr'][ind].index(1)+1)
    df['pos'][ind]=df['arr'][ind].index(1)+1

错误原因 #

你碰到KeyError: pos是因为在给df['pos'][ind]赋值前，根本没创建pos这一列。Pandas不允许直接对不存在的列做索引赋值，必须先初始化这个列。

另外你的代码还有两个问题：

• 循环只处理了前几行（df.head()的内容），没覆盖全部10000行数据
• 最后赋值的是数字，而你需要的是对应的分类标签（A/B/C/D）

正确实现方式 #

方法一：用idxmax高效处理（强烈推荐）

独热编码里每行值为1的列就是目标分类，Pandas的idxmax方法可以直接获取每行最大值（也就是1）所在的列名，一行代码搞定：

import pandas as pd

df = pd.read_csv("ABC.csv")
cols = ['A', 'B', 'C', 'D']
# 直接生成pos列
df['pos'] = df[cols].idxmax(axis=1)
# 只保留需要的列输出
result = df[['ids', 'pos']]
print(result.head())

方法二：修复你的原有代码

如果想保留自己的思路，需要先初始化pos列，再遍历所有行，同时把数字索引对应到分类标签：

import pandas as pd

df = pd.read_csv("ABC.csv")
cols = ['A', 'B', 'C', 'D']
# 先初始化pos列
df['pos'] = ''
df['arr'] = df[cols].values.tolist()

# 遍历所有行的索引
for ind in df.index:
    # 找到值为1的位置，对应cols里的标签
    pos_index = df['arr'][ind].index(1)
    # 用loc赋值避免警告和错误
    df.loc[ind, 'pos'] = cols[pos_index]

# 查看结果
print(df[['ids', 'pos']].head())

注意：循环处理Pandas数据框效率很低，面对10000行数据的话，优先用方法一。

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