如何基于多列特定条件创建以列名为值的新列？

给定如下DataFrame：

import numpy as np
import pandas as pd

df = pd.DataFrame({"ID":[1, 1, 2, 3, 3, 4],
                   "CLASS A":[np.nan, 1, np.nan, np.nan, np.nan, np.nan],
                   "CLASS B":[np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
                   "CLASS C":[np.nan, np.nan, np.nan, 1, np.nan, np.nan],
                   "DEPTH":[12, 31, 45, 66, 32, 46]})

需要新增列FILTER，仅在CLASS A、CLASS B、CLASS C列中值为1的行填入对应列名，其余行填充NaN，最终结果如下：

ID  CLASS A  CLASS B  CLASS C  DEPTH     FILTER
0   1      NaN      NaN      NaN     12     NaN
1   1      1.0      NaN      NaN     31     CLASS A
2   2      NaN      NaN      NaN     45     NaN
3   3      NaN      NaN      1.0     66     CLASS C
4   3      NaN      NaN      NaN     32     NaN
5   4      NaN      NaN      NaN     46     NaN

方法一：利用idxmax结合条件过滤 #

这是效率较高的向量化方法，适合大数据集：

class_cols = ['CLASS A', 'CLASS B', 'CLASS C']
# 先获取每行中值为1的列名，再过滤掉无1的行
df['FILTER'] = df[class_cols].idxmax(axis=1).where(df[class_cols].eq(1).any(axis=1))

逻辑说明：df[class_cols].eq(1).any(axis=1)判断每行是否存在值为1的分类列；idxmax会返回每行中第一个最大值的列名（这里1是唯一非NaN的有效值，因此能精准定位到目标列）；最后用where将不满足条件的行设为NaN。

方法二：通过melt重塑数据后匹配 #

适合需要保留中间过程的场景：

class_cols = ['CLASS A', 'CLASS B', 'CLASS C']
# 将分类列转为长格式，筛选出值为1的记录
melted_data = df.melt(id_vars=['ID', 'DEPTH'], value_vars=class_cols, var_name='FILTER')
melted_data = melted_data[melted_data['value'] == 1]
# 合并回原DataFrame，无匹配的自动填充NaN
df = df.merge(melted_data[['ID', 'DEPTH', 'FILTER']], on=['ID', 'DEPTH'], how='left')

逻辑说明：melt把多列分类数据转为键值对形式，筛选出值为1的行后，通过merge将对应列名匹配回原数据的对应行。

方法三：逐行遍历（小数据集适用） #

代码直观但效率较低，仅推荐用于小规模数据：

class_cols = ['CLASS A', 'CLASS B', 'CLASS C']
df['FILTER'] = np.nan

for idx, row in df.iterrows():
    for col in class_cols:
        if row[col] == 1:
            df.loc[idx, 'FILTER'] = col
            break

逻辑说明：遍历每一行，逐个检查分类列是否为1，找到匹配项后赋值对应列名并跳出循环，未找到则保持NaN。

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