CODEWITHSUNDEEP
pythonpandas
Abraham Mathew
Abraham Mathew

Reputation: 51

Select the max value for each group

So I have a pandas data frame with multiple columns and a id column.

df = pd.DataFrame(np.random.randn(6,4), columns=list('ABCD'))
df['id'] = ['CA', 'CA', 'CA', 'FL', 'FL', 'FL']
df['technique'] = ['one', 'two', 'three', 'one', 'two', 'three']
df

I want to group by the id column and select the row which had the highest probability. So it could look like this.

id   highest_prob   technique
CA   B               three 
FL   C               one

I tried something like this, but that would only get me half of the way. 

df.groupby('id', as_index=False)[['A','B','C','D']].max()

Anyone have suggestions on how I can get the desired result

Upvotes: 1

Views: 84

Answers (1)

cs95
cs95

Reputation: 402253

Setup

np.random.seed(0)  # Add seed to reproduce results. 
df = pd.DataFrame(np.random.randn(6,4), columns=list('ABCD'))
df['id'] = ['CA', 'CA', 'CA', 'FL', 'FL', 'FL']
df['technique'] = ['one', 'two', 'three', 'one', 'two', 'three']

You could melt, sort with sort_values, and drop duplicates using drop_duplicates:

(df.melt(['id', 'technique'])
   .sort_values(['id', 'value'], ascending=[True, False])
   .drop_duplicates('id')
   .drop('value', 1)
   .reset_index(drop=True)
   .rename({'variable': 'highest_prob'}, axis=1))

   id technique highest_prob
0  CA       one            D
1  FL       two            A

Another solution is to use melt and groupby:

v = df.melt(['id', 'technique'])
(v.iloc[v.groupby('id').value.idxmax()]
  .drop('value', 1)
  .reset_index(drop=True)
  .rename({'variable': 'highest_prob'}, axis=1))

   id technique highest_prob
0  CA       one            D
1  FL       two            A

Upvotes: 2

Related Questions