Find column where the max values but return nothing if all values are zero occurs

user_12749988 0
Viewed 375

The Problem

I want to find the column name (row by row) in a dataframe where the max value occurs but if all values are zero, in the row, I want it to return an empty string or NaN. I can get the column where the max value occurs using idxmax but using this returns the first column name if all values are the same. I tried using idxmax in a df.apply lambda but I could never get it to work.

df = pd.DataFrame({'cat1':[1,3,0], 'cat2':[2,0,0], 'cat3':[2,1,0], 'issues':[3,2,0]})

  issue1  issue2  issue3    issues
0   1       2        2        3
1   3       0        1        2
2   0       0        0        0

So the issues column is just a count (not sum) of non-zero values in the issue1, issue2 and issue3 columns.

Desired output:

  issue1  issue2  issue3    issues  top_issue
0   1       2        2        3      issue2
1   3       0        1        2      issue1
2   0       0        0        0       NaN

I want the top issue column like above where it gives me the column name with the highest value in it (first occurrence is fine) but when there are no issues (in row 3) I want a Null or empty value.

Things I Tried

1. Using idxmax

I was able to get the top_issue using this:

issue_cols = ['issue1','issue2','issue3']
df['top_issue'] = df[issue_cols].idxmax(axis=1)

But when all values in the issue columns are zero it gives me the first issue column.

Output:

  issue1  issue2  issue3    issues  top_issue
0   1       2        2        3      issue2
1   3       0        1        2      issue1
2   0       0        0        0      issue1

2. Using idxmax with apply

I tried using apply - lambda but I kept getting errors with idxmax and was not able to figure it out.

df['top_issue'] = df.apply(lambda row: None if row['issues']==0 else row[issue_cols].idxmax(axis=1),axis=1)

Which gave me a ValueError

ValueError: axis must be fewer than the number of dimensions (1)

Changed axis to zero which gave a TypeError

TypeError: reduction operation 'argmax' not allowed for this dtype

3. Using nlargest

df['top_issue'] = df.apply(lambda row: None if row['issues']==0 else row[issue_cols].nlargest(1).index.tolist(),axis=1)

This resulted in the following error

TypeError: Cannot use method 'nlargest' with dtype object

So then I had to double check that all my values were not of type object which they werent so I didn't know where to go from there.

pandas dataframe python
2 Answers

Use DataFrame.apply on the columns that match the pattern issue\d+, on axis=1 then take nan if maximum is zero else take the maximum index on that axis.

result=(df.assign(top_issue=df[df.columns[df.columns.str.match('issue\d+')]]
        .apply(lambda x: float('nan') if x.max()==0
        else x.idxmax(), axis=1)
                   )
         )

OUTPUT:

   issue1  issue2  issue3  issues top_issue
0       1       2       2       3    issue2
1       3       0       1       2    issue1
2       0       0       0       0       NaN
user_91363480

You can first replace 0 with NaN and then apply idxmax to get you desired output.

issue_cols = ['issue1','issue2','issue3']
df['top_issue'] = df[issue_cols].replace(0, np.nan).idxmax(axis=1)

Output:

     issue1  issue2  issue3  issues top_issue
0       1       2       2       3    issue2
1       3       0       1       2    issue1
2       0       0       0       0       NaN
user_134879150
Related