Search code examples
pythonpandasgroupingassign

How to use assign function after a groupby in Pandas Dataframe


I'm curious if we are able to use the assign function after a groupby in Python. I don't want to do this by saving the grouped data and then use assign.

Let's say we have:

import pandas as pd

data=pd.DataFrame({'Group':['A','A','A','B','B','B'],'One':[1,12,4,5,4,3],'Two':[5,3,4,5,2,7,]})

  Group  One  Two
0     A    1    5
1     A   12    3
2     A    4    4
3     B    5    5
4     B    4    2
5     B    3    7

The following code does not work because the data used inside the assign function is not the grouped one:

data.groupby('Group').sum().assign(one_two=data['One']/data['Two'])

The expected output:

       One  Two  one_two
Group                   
A       17   12      1.416
B       12   14      0.857

Again, I don't want to assign the grouped data to and new data frame and then use the assign. I want to do this in the same line of code.


Solution

  • Using DataFrame.eval

    data.groupby('Group').sum().eval('one_two = One / Two')
    

           One  Two   one_two
    Group
    A       17   12  1.416667
    B       12   14  0.857143