Pandas查找行内多列首个指定值并返回对应列名/索引的方法

可行解决方案 #

核心思路是先构造布尔匹配矩阵，再通过取最大值索引的方式拿到每行第一个匹配位置的列名，代码如下：

import pandas as pd
import numpy as np

# 初始化示例DataFrame
df = pd.DataFrame({"A":[0,0],"B":[0,1],"C":[1,2],"D":[2,2]})
# 定义要匹配的目标值
target = 2

# 生成布尔矩阵：等于目标值的位置为True，其余为False
mask = df.eq(target)
# 取每行第一个True对应的列名，无匹配时返回NaN
df["TAG"] = mask.idxmax(axis=1).where(mask.any(axis=1), np.nan)

如果你需要取列索引而非列名，可使用如下写法：

# 取列索引，无匹配时默认返回-1
df["TAG_IDX"] = mask.to_numpy().argmax(axis=1).where(mask.any(axis=1), -1)

原理解释 #

• df.eq(target)会生成和原DataFrame结构完全一致的布尔矩阵，布尔值True等价于数值1，False等价于数值0
• idxmax(axis=1)会返回每行最大值对应的列名，默认会取第一个出现最大值的位置，正好匹配「找第一个等于目标值的列」的需求
• 末尾追加的where是为了兼容无匹配的场景，避免某行没有目标值时错误返回第一个列名
• 你之前使用的np.where是逐元素判断返回值，生成的结果和原DataFrame维度一致，直接赋值给单列必然维度不匹配，所以无法达到预期效果

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