CODEWITHSUNDEEP
pythonpandasloopsdataframefor-loop
YilGuk Seo
YilGuk Seo

Reputation: 77

for loop and adding additional columns groupby pandas dataframe in Python

below code is my original way.

import pandas as pd
data = {'id':[1001,1001,1001,1001,1001,1001,1001,1001,1002,1002,1002,1002,1002,1002,1002,1002],
    'name':['Tom', 'Tom', 'Tom', 'Tom','Tom', 'Tom', 'Tom', 'Tom','Jack','Jack','Jack','Jack','Jack','Jack','Jack','Jack'],
    'team':['A','A', 'B', 'B', 'C','C', 'D', 'D','A','A', 'B', 'B', 'C','C', 'D', 'D',],
    'year':[2011,2011,2012,2012,2013,2013,2014,2014,2011,2011,2012,2012,2013,2013,2014,2014],
    'avg':[0.500,0.400,0.300,0.200,0.100,0.200,0.300,0.400,0.500,0.400,0.300,0.200,0.100,0.200,0.300,0.400]}

df = pd.DataFrame(data)

print (df)

team_names = [c for c in df['team'].value_counts().index]
team_names

for i in team_names:
    df[i+'_vs_avg_2011'] = df.loc[(df['team']==i)&(df['year']==2011)].groupby(['id','name'])['avg'].transform('mean')
    df[i+'_vs_avg_2012'] = df.loc[(df['team']==i)&(df['year']==2012)].groupby(['id','name'])['avg'].transform('mean')
    df[i+'_vs_avg_2013'] = df.loc[(df['team']==i)&(df['year']==2013)].groupby(['id','name'])['avg'].transform('mean')
    df[i+'_vs_avg_2014'] = df.loc[(df['team']==i)&(df['year']==2014)].groupby(['id','name'])['avg'].transform('mean')
    print(i)

for the loop part I tried 

years_from_to = [str(i).zfill(2) for i in range(2011,2014)]
years_from_to

for i,j in team_names, years_from_to:
    df[i+'_vs_avg_'+j] = df.loc[(df['team']==i)&(df['year']==j)].groupby(['id','name'])['avg'].transform('mean')
    print(i)

ValueError: too many values to unpack (expected 2)

Is there a way to simplify this or fix this code?

Upvotes: 3

Views: 406

Answers (1)

jezrael
jezrael

Reputation: 862691

I think you can use DataFrame.pivot_table instaed loops with flattening columns in MultiIndex and then DataFrame.join to original DataFrame:

df1 = df.pivot_table(index=['id','name'],columns=['team','year'],values='avg', aggfunc='mean')
df1.columns = [f'{a}_vs_avg_{b}' for a, b in df1.columns]
print (df1)
           A_vs_avg_2011  B_vs_avg_2012  C_vs_avg_2013  D_vs_avg_2014
id   name                                                            
1001 Tom            0.45           0.25           0.15           0.35
1002 Jack           0.45           0.25           0.15           0.35

df = df.join(df1, on=['id','name'])
print (df)

Upvotes: 4

Related Questions