CODEWITHSUNDEEP
pythonpandasdataframepivot
enixon4
enixon4

Reputation: 51

Pandas unstack with multiindex columns

I have a pandas dataframe, which can be created with:

pd.DataFrame([[1,'a','green'],[2,'b','blue'],[2,'b','green'],[1,'e','green'],[2,'b','blue']], columns  = ['sales','product','color'], index = ['01-01-2020','01-01-2020','01-02-2020','01-03-2020','01-04-2020'])

and looks like: Original Dataframe

I would like to unstack the dataframe with the 'color' feature and create a multiindex by product of [green,blue],[sales,product] with the already existing columns as the second level of the column multiindex. The index of the dataframe is a date. The resultant dataframe that I would like can be created with the code:

pd.DataFrame([[1,'a',2,'b'],[2,'b',np.nan,np.nan],[1,'e',np.nan,np.nan],[np.nan,np.nan,2,'b']],columns = pd.MultiIndex.from_product([['green','blue'],['sales','product']]), index = ['01-01-2020','01-02-2020','01-03-2020','01-04-2020'])

and looks like: Unstacked Dataframe

Please note that the resultant dataframe will be shorter than the original due to the common date indices.

For the life of me, I have been unable to figure out how to pivot/unstack correctly to figure this out. I am trying to apply this to a very large dataframe, so performance will be key for me. Many thanks for any and all help!

Upvotes: 3

Views: 5498

Answers (1)

Scott Boston
Scott Boston

Reputation: 153460

Try this:

df.set_index('color', append=True).unstack().swaplevel(0, 1, axis=1).sort_index(axis=1)

Output:

color         blue         green      
           product sales product sales
01-01-2020       b   2.0       a   1.0
01-02-2020     NaN   NaN       b   2.0
01-03-2020     NaN   NaN       e   1.0
01-04-2020       b   2.0     NaN   NaN

Details:

  • Add 'color' to your existing index with append=True
  • Unstack the inner most index level, 'color' to add it to columns
  • Swap the multiindex column header levels and sort

As, @QuangHoang states:

df.set_index('color', append=True).stack().unstack([1,2])

Which is much faster,

4.13 ms ± 274 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

versus

2.78 ms ± 44.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Upvotes: 8

Related Questions