Pandas如何将指定列中值为1的位置对应列名存入新增h列

首先给出测试基础代码：

import pandas as pd
dfz = pd.DataFrame({'A' : [1,0,0,1,0,0],
                    'B' : [1,0,0,1,0,1],
                    'C' : [1,0,0,1,3,1],
                    'D' : [1,0,0,1,0,0]})

需求规则 #

遍历A、B、C、D四列，对应行单元格值为1时，h列对应位置存储该单元格所属列名。

实现代码 #

场景1：单行存在多个符合条件的列名时，全部拼接展示 #

dfz['h'] = dfz.loc[:, 'A':'D'].eq(1).dot(dfz.columns + ',').str.rstrip(',')

场景2：仅返回第一个符合条件的列名 #

# 先匹配值为1的单元格，取第一个匹配列名，无匹配时返回空值
match_mask = dfz.loc[:, 'A':'D'].eq(1)
dfz['h'] = match_mask.idxmax(axis=1).where(match_mask.any(axis=1), '')

输出效果（以场景1为例） #

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