Changing column names in python pivot-CodePudding

I have a table like given below:

ID       MONTH_ID       CATEGORY     Disease_COUNT
551     2021-11           NON_MET       1
580     2022-03            MET          1
580     2022-01            MET          1

I have a created pivot on above table so that unique values present in CATEGORY column becomes new column and get the count as per disease_count as given in below table:

ID      MONTH_ID    COUNT_MET  COUNT_NON_MET
551     2021-11       null        1
580     2022-03        1         null
580     2022-01        1         null

I used below given code to create above pivot:

df2 = pd.pivot_table(df1,index=['ID','MONTH_ID'],
                        columns=['CATEGORY'],
                         values=["Disease_COUNT"]
                        ).reset_index()
cols=  ['BH_ID', 'MONTH_ID']
cols.extend([x for x in df2.columns.get_level_values(1) if x != ''])
df2.columns = cols

df1 is the first data frame given above I need to modify above given code in a way that I get following columns as ID, MONTH_ID, Disease_COUNT_MET and Disease_COUNT_NON_MET i.e I need Disease_COUNT_MET instead of COUNT_MET and Disease_COUNT_MET instead COUNT_NON_MET.

I need to modify code dynamically I can't use rename function.

Given below is the final output I need.

ID          MONTH_ID    Disease_COUNT_MET   Disease_COUNT_NON_MET
551          2021-11         null                    1
580          2022-03          1                    null
580          2022-01          1                    null

Please help

CodePudding user response：

You can flatten the MultiIndex:

df2 = (df1.pivot_table(index=['ID','MONTH_ID'],
                       columns=['CATEGORY'],
                       values=["Disease_COUNT"])
          .pipe(lambda d: d.set_axis(d.columns.map('_'.join), axis=1))
          .reset_index()
      )

df2

output:

    ID MONTH_ID  Disease_COUNT_MET  Disease_COUNT_NON_MET
0  551  2021-11                NaN                    1.0
1  580  2022-01                1.0                    NaN
2  580  2022-03                1.0                    NaN