Pandas Rolling Sum For Multiply Values Separately

I have the following dataframe: a = pd.DataFrame({'unit': [2, 2, 3, 3, 3, 4, 4, 4, 5], 'date': [1, 2, 1, 2, 3, 1, 2, 3, 1], 'revenue': [1, 1, 3,

Solution 1:

IIUC, you can do rolling on groupby:

a['rolled_sum'] = (a.groupby('unit')
                    .rolling(2, on='date').sum()['revenue']
                    .groupby('unit').shift(1)
                    .to_numpy()
                  )

Output:

   unit  date  revenue  rolled_sum
0211NaN1221NaN2313NaN3325NaN43378.05416NaN6426NaN743212.08519NaN

Solution 2:

With your sorting you can mask where it shouldn't be set.

m = a.unit.eq(a.unit.shift()) & a.unit.eq(a.unit.shift(-1))
a['rolled_sum'] = (a.rolling(2, on='date').sum().shift(+1)['revenue']
                     .where(m.shift().fillna(False)))

---

   unit  date  revenue  rolled_sum
0211NaN1221NaN2313NaN3325NaN43378.05416NaN6426NaN743212.08519NaN