如何使用pandas高效根据多列非零值生成对应标签列

问题说明 #

现有结构如下的DataFrame：

a b c d e
1 0 0 0 0
0 2 0 0 0
3 0 0 0 0
0 0 0 1 0
0 0 1 0 0
0 0 0 0 1

需要新增label列，取值规则为：每行非零值对应的列名前加col前缀，预期输出效果如下：

a  b  c  d  e label
0  1  0  0  0  0  cola
1  0  2  0  0  0  colb
2  3  0  0  0  0  cola
3  0  0  0  1  0  cold
4  0  0  1  0  0  colc
5  0  0  0  0  1  cole

之前尝试的代码df['label'] = df['a'].apply(lambda x: 1 if x!=0)存在两个明显问题：

• 仅判断了a列的非零情况，未覆盖其余所有列
• 逻辑仅返回0/1标识，没有匹配列名、拼接要求的col前缀，无法得到目标结果

可直接运行的实现代码 #

针对当前每行仅存在1个非零值的场景，优先使用向量化运算实现（运行效率更高）：

import pandas as pd

# 构造原始数据集（可替换为你自己的数据读取逻辑）
df = pd.DataFrame(
    [[1,0,0,0,0],
     [0,2,0,0,0],
     [3,0,0,0,0],
     [0,0,0,1,0],
     [0,0,1,0,0],
     [0,0,0,0,1]],
    columns=['a','b','c','d','e']
)

# 生成label列
df['label'] = 'col' + df.ne(0).idxmax(axis=1)

代码逻辑说明 #

• df.ne(0)：生成与原DataFrame维度一致的布尔矩阵，值非零的位置标记为True
• .idxmax(axis=1)：逐行返回第一个为True的位置对应的列名，刚好匹配单非零值的场景
• 最后通过字符串拼接，给列名加上要求的col前缀即可

如果后续存在一行多个非零值、需要取第一个非零列的场景，也可以用逐行apply的写法实现：

df['label'] = df.apply(lambda row: 'col' + row[row != 0].index[0], axis=1)

运行上述代码后，得到的结果和预期效果完全一致。

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