Software Mechanic Software Mechanic - 3 months ago 44
Python Question

How to convert pandas dataframe rows into columns, based on category?

I have a pandas data frame with a category variable and some number variables. Something like this:

ls = [{'count':5, 'module':'payroll', 'id':2}, {'count': 53, 'module': 'general','id':2}, {'id': 5,'count': 35, 'module': 'tax'}, ]
df = pd.DataFrame.from_dict(ls)

The df looks like this:

count id module
0 5 2 payroll
1 53 2 general
2 35 5 tax

I want convert(transpose is the right word?) the module variables into columns and group by the id. So something like:

general_count id payroll_count tax_count
0 53.0 2 5.0 NaN
1 NaN 5 NaN 35.0

One approach to this would be to use apply:

df['payroll_count'] = x: df[][df.module=='payroll'])

However, this suffers from multiple drawbacks:

  1. Costly, and takes too much time

  2. Creates artifacts and empty
    dataframes that need to be cleaned up.

I sense there's a better way to achieve this with pandas groupby, but can't find a way to this same operation more efficiently. Please help.


You can use groupby by columns which first create new index and last column. then need aggreagate some way - I use mean, then convert one column DataFrame to Series by DataFrame.squeeze (then is not necessary remove top level of Multiindex in columns) and reshape by unstack. Last add_suffix to column name:

df = df.groupby(['id','module']).mean().squeeze().unstack().add_suffix('_count')
print (df)
module  general_count  payroll_count  tax_count
2                53.0            5.0        NaN
5                 NaN            NaN       35.0

Another solution with pivot, then need remove Multiindex from columns by list comprehension:

df = df.pivot(index='id', columns='module')
df.columns = ['_'.join((col[1], col[0])) for col in df.columns]
print (df)
    general_count  payroll_count  tax_count
2            53.0            5.0        NaN
5             NaN            NaN       35.0