Pandas groupby sum

Question

I have a dataframe as follows:

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

This is what I'm trying to get:

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

So I have this:

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

But how can I filter it such that the above only applies to rows where type=foo or bar or test?

Answer 1

I think you need groupby with unstack and then merge to original DataFrame:

df1 = df.groupby(['ref','type'])['amount'].sum().unstack().reset_index()
print (df1)
type  ref  bar  foo  test
0     001   55   15   190
1     002  155   55   260

df = pd.merge(df, df1, on='ref')
print (df)
    ref  type  amount  sums  bar  foo  test
0   001   foo      10    15   55   15   190
1   001   foo       5    15   55   15   190
2   001   bar      50    55   55   15   190
3   001   bar       5    55   55   15   190
4   001  test     100   190   55   15   190
5   001  test      90   190   55   15   190
6   002   foo      20    55  155   55   260
7   002   foo      35    55  155   55   260
8   002   bar      75   155  155   55   260
9   002   bar      80   155  155   55   260
10  002  test     150   260  155   55   260
11  002  test     110   260  155   55   260

Timings:

In [506]: %timeit (pd.merge(df, df.groupby(['ref','type'])['amount'].sum().unstack().reset_index(), on='ref'))
100 loops, best of 3: 3.4 ms per loop

In [507]: %timeit (pd.merge(df, pd.pivot_table(df, values='amount', index=['ref'], columns=['type'], aggfunc=np.sum), left_on='ref', right_index=True))
100 loops, best of 3: 4.99 ms per loop

Answer 2

A solution using pivot table :

>>> 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