Fixing column names and renaming them after grouping the dataframe by two columns

I have a dataframe:

{'ARTICLE_ID': {0: 111, 1: 111, 2: 222, 3: 222, 4: 222}, 'CITEDIN_ARTICLE_ID': {0: 11, 1: 11, 2: 11, 3: 22, 4: 22}, 'enrollment': {0: 10, 1: 10, 2: 10, 3: 10, 4: 10}, 'Trial_year': {0: 2017, 1: 2017, 2: 2017, 3: 2017, 4: 2017}, 'AUTHOR_ID': {0: 'aaa', 1: 'aaa', 2: 'aaa', 3: 'aaa', 4: 'aaa'}, 'AUTHOR_RANK': {0: 1, 1: 2, 2: 3, 3: 4, 4: 5}}

I am grouping it by two columns

df_grouped = df.groupby(['AUTHOR_ID', 'Trial_year']).agg({'ARTICLE_ID': "count", 
                                                      'enrollment': ["count", 'sum']}).reset_index()

As a result, I receive this dataframe, where column names have two levels

{('AUTHOR_ID', ''): {0: 'aaa'}, ('Trial_year', ''): {0: 2017}, ('ARTICLE_ID', 'count'): {0: 5}, ('enrollment', 'count'): {0: 5}, ('enrollment', 'sum'): {0: 50}}

My ideal output – the dataframe with one level of column names and renamed column names

`AUTHOR_ID`, `Trial_year`, `ARTICLE_ID_count`, `enrollment_count`, `enrollment_sum`
     

>Solution :

You can modify the columns:

df_grouped.columns = [f"{i}_{j}" if j!='' else i for i,j in df_grouped.columns]

or use NamedAgg from the beginning:

df_grouped = (df.groupby(['AUTHOR_ID', 'Trial_year'])
              .agg(ARTICLE_ID_count=('ARTICLE_ID', "count"), 
                   enrollment_count=('enrollment','count'),
                   enrollment_sum=('enrollment','sum')).reset_index())

You can also pass a dictionary to groupby.agg for a little concise code:

df_grouped = (df.groupby(['AUTHOR_ID', 'Trial_year'], as_index=False)
              .agg(**{'_'.join(pair): pair for pair in [('ARTICLE_ID', 'count'), 
                                                        ('enrollment','count'), 
                                                        ('enrollment','sum')]}))

Output:

  AUTHOR_ID  Trial_year  ARTICLE_ID_count  enrollment_count  enrollment_sum
0       aaa        2017                 5                 5              50