pandas row wise sum when when consecutive column value is less than a certain number

Question

I have a data frame like this,

df
col1     col2    col3
 A       34       1
 B       86       2
 A       53       21
 C       24       33
 B       21       2
 C       11       1

Now I want to add col1 and col2 values row wise where consecutive col3 values are less than 3, so the final data frame would look like,

 col1    col2
   A      120
   A       53
   C       24
   B       32

I could do this using a for loop and comparing it with the previous row, but the execution time will be huge, looking for some pandas shortcuts to do this most efficiently.

Answer 1

You can use cumsum to get consecutive blocks of value <=3:

s = df.col3.ge(3)

# print `s.cumsum()` and `s` to see details
df.groupby([s.cumsum(),s], as_index=False).agg({'col1':'first','col2':'sum'})

Output:

  col1  col2
0    A   120
1    A    53
2    B    32
3    C    24