Pandas: Group By and Conditional Sum and Add Back to Data Frame

user_11303509 0
Viewed 97

I have a data frame like below:

ID  Num  Letter  Count
1   17   D       1
1   12   D       2
1   13   D       3
2   17   D       4
2   12   A       5
2   16   D       1
3   16   D       1

The objective is to sum 'Count' value for each 'ID' when 'Num' is (17 or 12) and 'Letter' is 'D', and also add the calculation back to the original data frame in 'Total'.

Below is expected data frame:

ID  Num  Letter  Count Total
1   17   D       1     3   
1   12   D       2     3   
1   13   D       3     3 
2   17   D       4     4 
2   12   A       5     4 
2   16   D       1     4
3   16   D       1     0

Thanks in advance!

pandas python
1 Answers

Idea is replace non matched values to 0 in Series.where and then is used GroupBy.transform with sum:

mask = df['Num'].isin([17,12]) & df['Letter'].eq('D')
df['Total'] = df['Count'].where(mask, 0).groupby(df['ID']).transform('sum')
print (df)
   ID  Num Letter  Count  Total
0   1   17      D      1      3
1   1   12      D      2      3
2   1   13      D      3      3
3   2   17      D      4      4
4   2   12      A      5      4
5   2   16      D      1      4
6   3   16      D      1      0
user_29010020
Related