How to sort a multindex pandas dataframe on an overall group level followed by within the group?

Question

Suppose I have the following dataframe:

import pandas as pd

tuples = [('A', 'AA'), ('A', 'AB'), ('B', 'BA'), ('B', 'BB'), 
          ('C', 'CA'), ('C', 'CB')]
index = pd.MultiIndex.from_tuples(tuples, 
                                  names=['first_level', 'second_level'])

input_df = pd.DataFrame([100, 500, 200, 50, 3000, 10000], 
                         columns=['amount'], index=index)
input_df

                            amount
first_level  second_level   
A                AA            100
                 AB            500
B                BA            200
                 BB             50
C                CA           3000
                 CB          10000

What I want to do is to sort based on two criteria: (1) The total amount across first_level overall and then (2) By the amount within each second_level .

In other words I want something like this:

tuples = [('C', 'CB'), ('C', 'CA'), ('A', 'AB'), 
          ('A', 'AA'), ('B', 'BA'), ('B', 'BB'), ]
index = pd.MultiIndex.from_tuples(tuples, 
                                  names=['first_level', 'second_level'])

output_df = pd.DataFrame([10000, 3000, 500, 100, 200, 50], 
                         columns=['amount'], index=index)
output_df


                                amount
first_level     second_level    
C                  CB            10000
                   CA             3000
A                  AB              500
                   AA              100
B                  BA              200
                   BB               50

As you can see group C has the largest amount (13000), followed by group A (600), and then group B (250). Within each group, the second_level is organized based on amount.

I have figured out one way of doing but it feels overly complicated as it requires aggregations, joins, and playing around with the index:

overall_group_amounts = input_df.groupby(['first_level']) \
    .sum() \
    .rename(columns={'amount': 'overall_amounts'})

pd.merge(overall_group_amounts, input_df.reset_index('second_level'), on='first_level') \
    .sort_values(['overall_amounts', 'amount', 'first_level'], ascending=[False, False, True]) \
    .drop('overall_amounts', axis='columns') \
    .set_index('second_level', append=True)

My question is: is there a better way of solving this problem?

Answer 1

You can create a temp sort key by summing each group then sort by the key and amount at the same time:

(
    df.assign(sk=df.groupby(level=0).amount.transform(sum))
    .sort_values(by=['sk','amount'], ascending=False)
    .drop('sk', 1)
)

                                amount
first_level     second_level    
C                  CB            10000
                   CA             3000
A                  AB              500
                   AA              100
B                  BA              200
                   BB               50