2
\$\begingroup\$

I have a task that is quite simple: I need to get the sum of values in the categories of each ID, and keep the category with the highest sum:

id    category   value
1       A         10 
1       A         15
1       B         13
2       A         80

So, in this case the sum of value for each category-id pair would be:

id    category   value
1       A         25 
1       B         13
2       A         80 

And then the maximum for id == 1 is 25 and for the other is 80, so the final dataframe is:

id    category   value
1       A         25 
2       A         80 

I achieved this like this:

(df.groupby(['id', 'category'])['value']
 .sum().reset_index().sort_values(by=['id', 'value'])
 .drop_duplicates(['id'], keep='last'))

I feel this can be done in lesser steps, maybe with some transform, but I can't find a better way. Any ideas?

Thanks

\$\endgroup\$
1
\$\begingroup\$

You could also do the following. Compute the first grouping by id and category, sum up the values:

y = df.groupby(["id","category"])["value"].sum()

Afterwards, grab the best category according to your definition:

y.groupby("category").sum().nlargest(1)

Combining these, so that we get the full job done:

y = df.groupby(["id","category"])["value"].sum()
cat = y.groupby("category").sum().nlargest(1).index
y.loc[:,cat]
\$\endgroup\$
1
  • \$\begingroup\$ I believe the original target was a max-sum category per id, rather than a single category overall? So y.loc[y.groupby('id').idxmax()] \$\endgroup\$ – GeoMatt22 Oct 15 '20 at 15:14

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.