Pandas: Get column index/label of a specific single value in a row from among multiple columns of mutually exclusive choice

user_16401663 0
Viewed 55

I have a pandas data frame like this, where the columns are the age groups of a user, and the number 1 indicates if the user is in that user group.

   21_30   31_40    40_49     50_59
0   0         0        0         1
1   0         1        0         0
2   0         1        0         0
3   1         0        0         0
4   0         0        1         0

What I'd like to do is gather all that information into one column, and transform the number 1 to a string that indicates the age group.

    age_group
0      50_59
1      31_40
2      31_40
3      21_30
4      40_49

How could I best approach this issue? Any advice would be helpful. Thanks in advance!

pandas python
2 Answers

You can check for the column values equal 1 across the row and get the column index corresponding to True (for entries of 1) by .idxmax() (with axis=1 for getting column index), as follows:

df['age_group'] = df.eq(1).idxmax(axis=1)

Result:

print(df)

   21_30  31_40  40_49  50_59 age_group
0      0      0      0      1     50_59
1      0      1      0      0     31_40
2      0      1      0      0     31_40
3      1      0      0      0     21_30
4      0      0      1      0     40_49
user_150706970

Try:

df['age_group'] = df.dot(df.columns)

Output:

0    50_59
1    31_40
2    31_40
3    21_30
4    40_49
dtype: object
user_63615310
Related