Adding sum of batches of rows in dataframe

Viewed 113

If i have a DataFrame, how can I compute the sum of a columns N rows based on the value of another column being consistent(?)

Example of dataframe:

      C1  C2  C3
0    400   F  31
1     10   F  32
2    300   F  33
3    100  Kn  29
4   3000  Kn  28
5    200  Kn  26
6     10   F  30
7   5000   F  34
8  30000  Kn  28
9  30000  Kn  26

Now the complicated part, I would need the sum of the C1 rows with the same C2 name until the C2 name changes. So, first the sum of three first F, then sum of following 3 Kn, then sum of 2 F etc.

Example of example output:

      C1  C2  C3    sum
0    400   F  31    
1     10   F  32    
2    300   F  33    710
3    100  Kn  29   
4   3000  Kn  28   
5    200  Kn  26   3300
6     10   F  30   
7   5000   F  34   5010
8  30000  Kn  28  
9  30000  Kn  26  60000

I can retrieve all rows based on C2 value with df.loc[df['C2'] == 'F'], but then I can only get sum of all F rows, which is not what I want.

How can i get the sum of each F value until Kn appears, then get sum of Kn value until F appears, etc.

I'm having a hard time constructing this question to make sense, feel free to come with ideas on how I can improve the phrasing.

2 Answers

We can use cumsum with shift to create our custom groups based on changing calues in C2. Then we use these groups in GroupBy.sum to get our desired column:

grps = df['C2'].ne(df['C2'].shift()).cumsum()
df['sum'] = df.groupby(grps)['C1'].transform('sum')
      C1  C2  C3    sum
0    400   F  31    710
1     10   F  32    710
2    300   F  33    710
3    100  Kn  29   3300
4   3000  Kn  28   3300
5    200  Kn  26   3300
6     10   F  30   5010
7   5000   F  34   5010
8  30000  Kn  28  60000
9  30000  Kn  26  60000

A small addition to @Erfan's elegant answer. To remove the sum values except the last entry in a group, you can do this after groupby is done:

df.loc[~df['C2'].ne(df['C2'].shift(-1)), "sum"] = np.nan

df
Out[30]: 
      C1  C2  C3      sum
1    400   F  31      NaN
2     10   F  32      NaN
3    300   F  33    710.0
4    100  Kn  29      NaN
5   3000  Kn  28      NaN
6    200  Kn  26   3300.0
7     10   F  30      NaN
8   5000   F  34   5010.0
9  30000  Kn  28      NaN
9  30000  Kn  26  60000.0

where .shift(-1) shifts the series backward by 1.

Related