问题
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, 5, 7, 6, 6, 2, 9]})
Pandas rolling.sum with window = 2:
a['rolled_sum'] = a.rolling(2, on='date').sum().shift(+1)['revenue']
computes this sum row by row:
adunit date revenue rolled_sum
0 2 1 1 NaN
1 2 2 1 NaN
2 3 1 3 2.0
3 3 2 5 4.0
4 3 3 7 8.0
5 4 1 6 12.0
6 4 2 6 13.0
7 4 3 2 12.0
8 5 1 9 8.0
I would like to have this rolling sum computed for each unit separately:
adunit date revenue rolled_sum
0 2 1 1 NaN
1 2 2 1 NaN
2 3 1 3 NaN
3 3 2 5 NaN
4 3 3 7 8.0
5 4 1 6 NaN
6 4 2 6 NaN
7 4 3 2 12.0
8 5 1 9 NaN
In other words: rolling sum should be performed for each unit separately. In my original dataset I have hundreds of units, and want to perform a rolling sum day-by-day for each of them.
Any ideas?
Many Thanks in advance :)
Andy
回答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
0 2 1 1 NaN
1 2 2 1 NaN
2 3 1 3 NaN
3 3 2 5 NaN
4 3 3 7 8.0
5 4 1 6 NaN
6 4 2 6 NaN
7 4 3 2 12.0
8 5 1 9 NaN
回答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
0 2 1 1 NaN
1 2 2 1 NaN
2 3 1 3 NaN
3 3 2 5 NaN
4 3 3 7 8.0
5 4 1 6 NaN
6 4 2 6 NaN
7 4 3 2 12.0
8 5 1 9 NaN
来源：https://stackoverflow.com/questions/58510308/pandas-rolling-sum-for-multiply-values-separately