Сгруппировать мультииндекс и получить сумму для первого индекса

Question

У меня есть данные, которые я хочу сгруппировать и суммировать по мультииндексу, но затем также иметь сумму только первой части мультииндекса.Таким образом, идея заключается в том, что индексом будет мультииндекс, и тогда я смог увидеть как размер этого, так и размер только его первой группы.Например, у меня есть такой фрейм данных:

d = {'attr_1': ['A','A','A','B','B'], 'attr_2': ['a','a','b','a','b'],'val': [1,2,3,4,5]}
df = pd.DataFrame(d)
df

    a_1 a_2 val
0   A   a   1
1   A   a   2
2   A   b   3
3   B   a   4
4   B   b   5

Я хотел бы сделать следующее:

        val sum_a1
a_1 a_2 
A   a   3   6
A   b   3   6
B   a   4   9
B   b   5   9

Я ищу способ сделать это чисто.Мой текущий подход состоит в том, чтобы создать новый фрейм данных, превратить его в словарь, сбросить индекс и отобразить словарь, но, похоже, для этого должен быть более простой и более питонический способ.

a = df.groupby(['a_1','a_2']).sum()
b = df.groupby(['a_1']).sum()
my_dict = b.to_dict()['val']
a = a.reset_index()
a['sum_a1'] = a.a_1.map(my_dict)
a
    a_1 a_2 val sum_a1
0   A   a   3   6
1   A   b   3   6
2   B   a   4   9
3   B   b   5   9

Любая помощь будет принята с благодарностью.Извиняюсь, если это повторяющийся вопрос, который я искал и не смог найти ничего похожего.

Answer 1

Вы ищете transform

a = df.groupby(['attr_1','attr_2']).sum()
a['sum_a1']=a.groupby(level=0)['val'].transform('sum')
a
Out[14]: 
               val  sum_a1
attr_1 attr_2             
A      a         3       6
       b         3       6
B      a         4       9
       b         5       9

Answer 2

С set_index и sum, позволяющими join выровнять индексы.

df.set_index(['attr_1', 'attr_2']).pipe(
    lambda d: d.sum(level=[0, 1]).join(d.val.rename('sum_a1').sum(level=0))
)

               val  sum_a1
attr_1 attr_2             
A      a         3       6
       b         3       6
B      a         4       9
       b         5       9