Pandas find the first occurrence of a specific value in a row within multiple columns and return column index

Question

For a dataframe:
df = pd.DataFrame({"A":[0,0],"B":[0,1],"C":[1,2],"D":[2,2]})

How to obtain the column name or column index when the value is 2 or a certain value

and put it in a new column at df, say df["TAG"]
df = pd.DataFrame({"A":[0,0],"B":[0,1],"C":[1,2],"D":[2,2],"TAG":[D,C]})

i tried
df["TAG"]=np.where(df[cols]>=2,df.columns,'')

where [cols] is the list of df columns

So far i can only find how to find row index when matching a value in Pandas

In excel we can do some approach using MATCH(TRUE,INDEX($A:$D>=2,0),) and apply to multiple rows

Any help or hints are appreciated

Thank you so much in advance

user_8708364 · Answer

Try idxmax:

>>> df['TAG'] = df.ge(2).T.idxmax()
>>> df
   A  B  C  D TAG
0  0  0  1  2   D
1  0  1  2  2   C
>>>