Conditional row amount in pandas groupby expression

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=2it 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 groupbyand 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)
0
source share
1

pd.rolling_sum window=2, shift NaNs 0

In [114]: df['l2'] = (df.groupby('product')['sales']
                       .apply(lambda x: pd.rolling_sum(x, window=2, min_periods=0)
                       .shift()
                       .fillna(0)))
In [115]: df
Out[115]:
   product  sales  week    l2
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
+2

Source: https://habr.com/ru/post/1613477/


All Articles