熊猫组合总和

Question

我有一个数据帧如下:

ref, type, amount
001, foo, 10
001, foo, 5
001, bar, 50
001, bar, 5
001, test, 100
001, test, 90
002, foo, 20
002, foo, 35
002, bar, 75
002, bar, 80
002, test, 150
002, test, 110

这就是我想要得到的:

ref, type, amount, foo, bar, test
001, foo, 10, 15, 55, 190
001, foo, 5, 15, 55, 190
001, bar, 50, 15, 55, 190
001, bar, 5, 15, 55, 190
001, test, 100, 15, 55, 190
001, test, 90, 15, 55, 190
002, foo, 20, 55, 155, 260
002, foo, 35, 55, 155, 260
002, bar, 75, 55, 155, 260
002, bar, 80, 55, 155, 260
002, test, 150, 55, 155, 260
002, test, 110, 55, 155, 260

所以我有这个:

df.groupby('ref')['amount'].transform(sum)

但我怎么能过滤它使得上面只适用于行,其中type=foo或bar或test？

Answer 1

使用溶液透视表:

>>> b = pd.pivot_table(df, values='amount', index=['ref'], columns=['type'], aggfunc=np.sum)
>>> b
type  bar  foo  test
ref
1      55   15   190
2     155   55   260

>>> pd.merge(df, b, left_on='ref', right_index=True)
    ref  type  amount  bar  foo  test
0     1   foo      10   55   15   190
1     1   foo       5   55   15   190
2     1   bar      50   55   15   190
3     1   bar       5   55   15   190
4     1  test     100   55   15   190
5     1  test      90   55   15   190
6     2   foo      20  155   55   260
7     2   foo      35  155   55   260
8     2   bar      75  155   55   260
9     2   bar      80  155   55   260
10    2  test     150  155   55   260
11    2  test     110  155   55   260