3

I working on a Python project that has a DataFrame like this:

data = {'AAA':  [3, 8, 2, 1],
        'BBB':  [5, 4, 7, 2],
        'CCC':  [2, 5, 6, 4]}
df = pd.DataFrame(data)

which leads to:

AAA BBB CCC
0 3 5 2
1 8 4 5
2 2 7 6
3 1 2 4

And the task consists of generating the following DataFrame:

AAA BBB CCC Role
0 3 5 2 BBB
1 8 4 5 AAA
2 2 7 6 BBB
3 1 2 4 CCC

Where "Role" column elements are the column headers that have the highest value in the row in which it is located.

Could you please help me by suggesting a code that solves this task?

Wagner André Yamada Vieira
  • 95
  • 4

1 Answers1

5

You could use the idxmax method on axis:

df['Role'] = df.idxmax(axis=1)

Output:

   AAA  BBB  CCC  Role
0    3    5    2  BBB
1    8    4    5  AAA
2    2    7    6  BBB
3    1    2    4  CCC