熊猫如何根据其他列中的值汇总一列的总和

Question

I am trying to sum values in a column by groupby on values in a second column, but meanwhile also considering values on a 3rd column, the df is like,

id    memo    amount   
 1    pos     1.0 
 1    pos     2.0
 1    neg     3.0
 2    pos     4.0
 2    pos     5.0
 2    neg     6.0
 2    neg     7.0

I want to group by id and sum amount , but each group, if memo is pos it is positive and neg for negative, eg when groupby 1 , the total amount is 0, since -1.0 - 2.0 + 3.0 = 0 .

If I do df.groupby('id')['amount'].sum() , it only considers id and amount column, I am wondering how to also take memo into account here.

so the result will look like,

id    memo    amount    total_amount   
 1    pos     1.0       0.0
 1    pos     2.0       0.0
 1    neg     3.0       0.0
 2    pos     4.0       -4.0
 2    pos     5.0       -4.0
 2    neg     6.0       -4.0
 2    neg     7.0       -4.0

Answer 1

Splitting the operation in two steps, you can achieve what you want through

df['temp'] = np.where(df.memo == 'pos', df.amount, -df.amount)
df['total_amount'] = df.groupby('id').temp.transform(sum)

Answer 2

Another fun way with mapping and multiplying ie

df['new'] = (df.set_index('id')['memo'].map({'pos':1,'neg':-1})*df['amount'].values)\
            .groupby(level=0).transform(sum).values

Output :

   id memo  amount  new
0   1  pos     1.0  0.0
1   1  pos     2.0  0.0
2   1  neg     3.0  0.0
3   2  pos     4.0 -4.0
4   2  pos     5.0 -4.0
5   2  neg     6.0 -4.0
6   2  neg     7.0 -4.0