Cumulative sum in pandas starting with a zero and ending with the sum of all but the last entry respecting groups

user_15039319 0
Viewed 1118

In the dataframe below, I want to create a new column C which will be the cumulative sum of B for each group in the A column, but these sums must start at zero and only add the values until the penultimate entry for that group.

  A B
0 1 5
1 1 6
2 2 3
3 2 4
4 2 5
5 3 2
5 3 7
6 4 3

So, my result should be:

  A B C
0 1 5 0
1 1 6 5
2 2 3 0
3 2 4 3
4 2 5 7
5 3 2 0
5 3 7 2
6 4 3 0

(I think this question is really obvious but somehow I couldn't figure it out myself nor I could see it anywhere asked already.)

cumsum pandas-groupby pandas python
3 Answers

Try:

df["C"] = df.groupby("A")["B"].transform(
    lambda x: x.shift().fillna(0).cumsum().astype(int)
)
print(df)

Prints:

   A  B  C
0  1  5  0
1  1  6  5
2  2  3  0
3  2  4  3
4  2  5  7
5  3  2  0
5  3  7  2
6  4  3  0
user_100359850

Given your groups are already consecutive you can shift and use where to NaN the rows that cross groups. The shift also ensures the cumsum is with respect to previous rows and stats at 0, as you want. Then use the built-in groupby + cumsum to avoid the slower lambda.

s = df['A'].shift()

df['C'] = (df.shift()
             .where(df['A'].eq(s))
             .groupby('A')['B'].cumsum()
             .fillna(0, downcast='infer'))


#   A  B  C
#0  1  5  0
#1  1  6  5
#2  2  3  0
#3  2  4  3
#4  2  5  7
#5  3  2  0
#5  3  7  2
#6  4  3  0

If it turns out that the grouping rows aren't consecutive we can handle this too. The only small modification is that we need to stable sort based on the grouping columns (such that the order within group is preserved and the cumsum is correct) and then in the end we can sort the index to get back to the original ordering of the DataFrame, if that's important.

# Create DataFrame with non-consecutive groups, 
df = pd.concat([df[::2], df[1::2]], ignore_index=True)
#   A  B
#0  1  5
#1  2  3
#2  2  5
#3  3  7
#4  1  6
#5  2  4
#6  3  2
#7  4  3

df = df.sort_values('A', kind='mergesort')

s = df['A'].shift()

df['C'] = (df.shift()
             .where(df['A'].eq(s))
             .groupby('A')['B'].cumsum()
             .fillna(0, downcast='infer'))

df = df.sort_index()
#   A  B  C
#0  1  5  0
#1  2  3  0
#2  2  5  3  <- Previous row `1` has value 3
#3  3  7  0
#4  1  6  5  <- Previous row `0` has value 5
#5  2  4  8  <- Previous rows `1` and `2` has values 3 + 5 = 8
#6  3  2  7  <- Previous row `3` has value 7 
#7  4  3  0
user_43333590

Another option is to use .groupby() twice, as follows:

Take the DataFrameGroupBy.shift() value of B under A so that for each group of A, the first entry of B will be reset and become NaN for later .fillna() to 0.

Further grouped by A for GroupBy.cumsum() within the local sequence of A to get the desired output:

df['C'] = (df.groupby('A')['B'].shift()
             .groupby(df['A']).cumsum()
             .fillna(0, downcast='infer')
          )

This solution is vectorized as well as supporting non-contiguous groups too!

Result:

print(df)


   A  B  C
0  1  5  0
1  1  6  5
2  2  3  0
3  2  4  3
4  2  5  7
5  3  2  0
5  3  7  2
6  4  3  0

Edit

If you are to groupby more than one column and you got "KeyError", check whether your syntax is correctly entered, for example:

If you groupby 2 columns year and income, you can use:

df['C'] = (df.groupby(['year', 'income'])['B'].shift()
             .groupby([df['year'], df['income']]).cumsum()
             .fillna(0, downcast='infer')
          )

Pandas supports both syntax with or without quoting df passing parameter to .groupby(). However, for any groupby() that the entity being grouped is not df itself, we may not be able to use the abbreviated form to quote just the column labels only e.g. 'year', we need to use the full column name e.g. df['year'] instead.

user_150706970
Related