Pandas DataFrame добавляет сумму за идентификатор

Question

У меня есть df как это:

idx total
1   5
1   10
1   2
2   10
2   10
2   55

И я хотел бы создать еще один столбец, который суммирует и усредняет итоги по идентификатору, чтобы получить это:

idx total sum   avg
1   5     17    5.66
1   10    17    5.66
1   2     17    5.66
2   10    75    25
2   10    75    25
2   55    75    25

Answer 1

Вы можете использовать groupby с agg, затем merge:

df.merge(df.groupby('idx')['total']
           .agg(['sum','mean']), 
         left_on='idx', 
         right_index=True)

Выход:

   idx  total  sum       mean
0    1      5   17   5.666667
1    1     10   17   5.666667
2    1      2   17   5.666667
3    2     10   75  25.000000
4    2     10   75  25.000000
5    2     55   75  25.000000

Answer 2

Обычно мы должны использовать transform, но, поскольку transform не может принять список как функцию с groupby объектом, мы используем agg + reindex

df.groupby('idx').total.agg(['sum','mean']).reindex(df.idx).reset_index()
Out[64]: 
   idx  sum       mean
0    1   17   5.666667
1    1   17   5.666667
2    1   17   5.666667
3    2   75  25.000000
4    2   75  25.000000
5    2   75  25.000000

#df.groupby('idx').total.agg(['sum','mean']).reindex(df.idx).reset_index().assign(Todat=df.total.values)