Блог

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

 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

  

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

Вы ищете 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
  

С 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