pandas dataframe groupby: apply a function that involves the grouping index value

user_12528151 0
Viewed 672

Let me use this simple df as an example.

df = pd.DataFrame([[1,1,10],[1,1,30],[1,2,40],[2,3,50],[2,3,150],[2,4,100]],columns=["col_grp","col1","col2"])

col_grp col1    col2
0   1   1   10
1   1   1   30
2   1   2   40
3   2   3   50
4   2   3   150
5   2   4   100

I wish to groupby 'col_grp' and aggregate every other columns into a single value. For example, grp['col1'].apply(lambda x: x.sum()). But instead of a simple 'sum' function f(x)=sum(x), I wish to compute a complicated function which takes not only on the pd.Series of 'col1' as its input, but also the grouping value, i.e., f(x,y). In this example, for the group with 'col_grp' value = 1, the y input is 1. The x input is the 'col1' Series [1,1,2] as usual. Assuming f(x,y) can be coded but its logic a bit complicated for other workaround methods, can I still utilize grp.apply(f)? If yes, how should I write the syntax. Thank you in advance.

I just found something like f(x1,x2) can be applied where x1 is the 'col1' series and x2 is the 'col2' series after grouping. f(x1,x2) returns a single value which is meant to be a statistic of each group. E.g.,

grp[['col1','col2']].apply(lambda x: x['col1'].sum() + x['col2'].median())

So this could be an effective way to solve my question. Just involve the grouping column into the x input. Like,

grp[['col_grp','col1','col2']].apply(lambda x: x['col_grp'].mean() + x['col1'].sum() + x['col2'].median())

And I can customize f(x) in anyway I like. I think I answered my own question. Welcome any further discussions:)

pandas-groupby pandas aggregate-functions dataframe python
1 Answers

Create a group (grp) with .groupby and create a calculated column using transform

import pandas as pd
df = pd.DataFrame([[1,1,10],[1,1,30],[1,2,40],[2,3,50],[2,3,150],[2,4,100]],columns=["col_grp","col1","col2"])
grp = df.groupby('col_grp')
df['new_col'] = grp['col1'].transform('sum') + grp['col2'].transform('median')
df
Out[1]: 
   col_grp  col1  col2  new_col
0        1     1    10       34
1        1     1    30       34
2        1     2    40       34
3        2     3    50      110
4        2     3   150      110
5        2     4   100      110

To make it your own function that you want to customize, you don't have to use apply, depending on what you are trying to do.

import pandas as pd
df = pd.DataFrame([[1,1,10],[1,1,30],[1,2,40],[2,3,50],[2,3,150],[2,4,100]],columns=["col_grp","col1","col2"])

def f(dataframe, col1, col2):
    grp = dataframe.groupby('col_grp')
    return grp[col1].transform('sum') + grp[col2].transform('median')

df['new_col'] = f(df, 'col1', 'col2')    
df
Out[2]: 
   col_grp  col1  col2  new_col
0        1     1    10       34
1        1     1    30       34
2        1     2    40       34
3        2     3    50      110
4        2     3   150      110
5        2     4   100      110
user_63667700
Related