Панды - группировка, агрегирование и масштабирование по сумме нескольких столбцов

Question

Предположим, у меня есть следующий DataFrame:

import pandas as pd
df = pd.DataFrame({'id': [1, 2, 2, 3, 3, 3], 'A': [2, 2, 3, 3, 5, 2], 'B': [1, 2, 1, 3, 2, 4]})
df
Out[253]: 
   id  A  B
0   1  2  1
1   2  2  2
2   2  3  1
3   3  3  3
4   3  5  2
5   3  2  4

Я бы хотел groupby 'id' и агрегировать с помощью функции sum над 'A', 'B'. Но я также хотел бы масштабировать A и B на сумму A+B (для каждого id), поэтому следующий результат будет выглядеть следующим образом:

   id         A         B
0   1  0.666667  0.333333
1   2  0.625000  0.375000
2   3  0.526316  0.473684

Теперь я могу сделать

res = df.groupby('id').agg('sum').reset_index()
scaler = res['A'] + res['B']
res['A'] /= scaler
res['B'] /= scaler
res
Out[275]: 
   id         A         B
0   1  0.666667  0.333333
1   2  0.625000  0.375000
2   3  0.526316  0.473684

Что довольно не элегантно . Есть ли способ поместить всю эту «скалярную» логику в функцию агрегирования? Или любой другой питонский и элегантный способ сделать это? Решения с участием numpy также приветствуются!

Answer 1

Вы можете сделать

In [584]: res = df.groupby('id').sum()

In [585]: res.div(res.sum(1), 0).reset_index()
Out[585]:
   id         A         B
0   1  0.666667  0.333333
1   2  0.625000  0.375000
2   3  0.526316  0.473684

Answer 2

Нет, вы не можете использовать функцию agg для масштабирования, потому что работаете с каждым столбцом отдельно.

Решение удалить reset_index для выравнивания в делении (div) Series, созданного sum:

res = df.groupby('id').sum()
res = res.div(res.sum(axis=1), axis=0).reset_index()
print (res)
   id         A         B
0   1  0.666667  0.333333
1   2  0.625000  0.375000
2   3  0.526316  0.473684

Подробности:

print (res.sum(axis=1))
id
1     3
2     8
3    19
dtype: int64

Answer 3

Вы можете использовать sum вдоль первой оси:

res = df.groupby('id').agg('sum')
res.div(res.sum(1), 0)

           A         B
id
1   0.666667  0.333333
2   0.625000  0.375000
3   0.526316  0.473684