Column Group Addition in Pandas

Question

I am working with a dataframe like this:

state current
  0     0
  0     0
  0     0
  0     0
  1     13
  1     7
  1     1
  1     12
  0     0
  0     0
  0     0
  1     8
  1     9

And I'd like to find the sums of the groups of currents where the state (and current) is not 0 and write this sum for all entries in the group. Ideally I would end up with something like this:

state current sum
  0     0      0
  0     0      0
  0     0      0
  0     0      0
  1     13     33
  1     7      33
  1     1      33
  1     12     33
  0     0      0
  0     0      0
  0     0      0
  1     8      17
  1     9      17

Is the best way to do that to divide up the dataframe into groups with values that arent 0, calculate the sums, then reassemble the dataframe?

Answer 1

When you groupby , first check the groupby key is 0 or not, then we do cumsum

df.groupby(df.state.eq(0).cumsum()).current.transform('sum')
0      0
1      0
2      0
3     33
4     33
5     33
6     33
7     33
8      0
9      0
10    17
11    17
12    17
Name: current, dtype: int64