I have a dataframe containing weekly sales for different products (a, b, c):
In[1]
df = pd.DataFrame({'product': list('aaaabbbbcccc'),
'week': [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4],
'sales': np.power(2, range(12))})
Out[1]
product sales week
0 a 1 1
1 a 2 2
2 a 4 3
3 a 8 4
4 b 16 1
5 b 32 2
6 b 64 3
7 b 128 4
8 c 256 1
9 c 512 2
10 c 1024 3
11 c 2048 4
I would like to create a new column containing cumulative sales for the last n weeks, grouped by product. For example. for n=2
it should be like last_2_weeks
:
product sales week last_2_weeks
0 a 1 1 0
1 a 2 2 1
2 a 4 3 3
3 a 8 4 6
4 b 16 1 0
5 b 32 2 16
6 b 64 3 48
7 b 128 4 96
8 c 256 1 0
9 c 512 2 256
10 c 1024 3 768
11 c 2048 4 1536
How can I efficiently calculate such a cumulative, notional sum in pandas? The solution should also work if there are more variables to group, for example. product and location.
I tried to create a new function and use groupby
and apply
, but this only works if the rows are sorted. He is also slow and ugly.
def last_n_weeks(x):
""" calculate sales of previous n weeks in aggregated data """
n = 2
cur_week = x['week'].iloc[0]
cur_prod = x['product'].iloc[0]
res = np.sum(df['sales'].loc[((df['product'] == cur_prod) &
(df['week'] >= cur_week-n) & (df['week'] < cur_week))])
return res
df['last_2_weeks'] = df.groupby(['product', 'week']).apply(last_n_weeks).reset_index(drop=True)