AJG519
AJG519

Reputation: 3379

Pandas get sorted index order for multiple columns

I have something like the following multi-index Pandas series where the values are indexed by Team, Year, and Gender.

>>> import pandas as pd
>>> import numpy as np
>>> multi_index=pd.MultiIndex.from_product([['Team A','Team B', 'Team C', 'Team D'],[2015,2016],['Male','Female']], names = ['Team','Year','Gender'])
>>> np.random.seed(0)
>>> df=pd.Series(index=multi_index, data=np.random.randint(1, 10, 16))
>>> df
>>> 
Team    Year  Gender
Team A  2015  Male      6
              Female    1
        2016  Male      4
              Female    4
Team B  2015  Male      8
              Female    4
        2016  Male      6
              Female    3
Team C  2015  Male      5
              Female    8
        2016  Male      7
              Female    9
Team D  2015  Male      9
              Female    2
        2016  Male      7
              Female    8

My goal is to get a dataframe of the team ranked order for each of the 4 Year / Gender combinations (Male 2015, Male 2016, Female 2015, and Female 2016).

My approach has been to first unstack the dataframe so that it is indexed by team...

>>> unstacked_df = df.unstack(['Year','Gender'])
>>> print unstacked_df
>>> 
>>> 
Year   2015        2016       
Gender Male Female Male Female
Team                          
Team A    6      1    4      4
Team B    8      4    6      3
Team C    5      8    7      9
Team D    9      2    7      8

And then create a dataframe from the index orders by looping through and sorting each of those 4 columns...

>>> team_orders = np.array([unstacked_df.sort_values(x).index.tolist() for x in unstacked_df.columns]).T
>>> result = pd.DataFrame(team_orders, columns=unstacked_df.columns)
>>> print result
Year      2015            2016        
Gender    Male  Female    Male  Female
0       Team C  Team A  Team A  Team B
1       Team A  Team D  Team B  Team A
2       Team B  Team B  Team C  Team D
3       Team D  Team C  Team D  Team C

Is there an easier / better approach that I'm missing?

Upvotes: 1

Views: 1066

Answers (1)

Randy
Randy

Reputation: 14847

Starting from your unstacked version, you can use .argsort() with .apply() to rank order each column and then just use that as a lookup against the index:

df.unstack([1,2]).apply(lambda x: x.index[x.argsort()]).reset_index(drop=True)

Year      2015            2016        
Gender    Male  Female    Male  Female
0       Team C  Team A  Team A  Team B
1       Team A  Team D  Team B  Team A
2       Team B  Team B  Team C  Team D
3       Team D  Team C  Team D  Team C

EDIT: Here's a little more info on why this works. With just the .argsort(), you get:

print df.unstack([1,2]).apply(lambda x: x.argsort())

Year   2015        2016       
Gender Male Female Male Female
Team                          
Team A    2      0    0      1
Team B    0      3    1      0
Team C    1      1    2      3
Team D    3      2    3      2

The lookup bit is essentially just doing the following for each column:

df.unstack([1,2]).index[[2,0,1,3]]

Index([u'Team C', u'Team A', u'Team B', u'Team D'], dtype='object', name=u'Team')

and the .reset_index() gets rid of the now-meaningless index labels.

Upvotes: 2

Related Questions