Update row value with column name if condition is met

Question

Im trying to do the following to a pandas dataframe:

Input data:

Char   A    B
Nan    0    1
Nan    1    0

Output:

Char   
B
A

As you can see I want to fill the value in 'Char' with a column name if the value of said column is 1. How should I approach this problem?

Answer 1

You can use

df['Char'] = df[['A','B']].idxmax(axis=1)
df

Out:

  Char  A  B
0    B  0  1
1    A  1  0