python - Pandas:按组对前 N 行求和

Question

我想对每组的前 N ​​期数据求和。我已经看到了如何单独执行每个操作(按组求和，或 sum prior N periods )，但无法找到一种干净的方法来同时执行这两项操作。

我目前正在做以下事情:

import pandas as pd

sample_data = {'user': ['a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b'],\
                'clicks': [0,1,2,3,4,5,6,7,8,9]}
df = pd.DataFrame(sample_data)
df['clicks.1'] = df.groupby(['user'])['clicks'].shift(1)
df['clicks.2'] = df.groupby(['user'])['clicks'].shift(2)
df['clicks.3'] = df.groupby(['user'])['clicks'].shift(3)
df['total_clicks_prior3'] = df[['clicks.1','clicks.2', 'clicks.3']].sum(axis=1)

我不想要 3 个中间滞后列，我只想要它们的总和，所以我想要的输出是:

>>> df[['clicks','user','total_clicks_prior3']]
   clicks user  total_clicks_prior3
0       0    a                  NaN
1       1    a                  0.0
2       2    a                  1.0
3       3    a                  3.0
4       4    a                  6.0
5       5    b                  NaN
6       6    b                  5.0
7       7    b                 11.0
8       8    b                 18.0
9       9    b                 21.0

注意:我显然可以在创建 3 列后删除它们，但考虑到我将创建多列不同数量的滞后周期，我觉得必须有更简单的方法。

Answer 1

这是groupby + rolling + shift

df.groupby('user')['clicks'].rolling(3, min_periods=1).sum().groupby(level=0).shift()

user
a     0     NaN
      1     0.0
      2     1.0
      3     3.0
      4     6.0
b     5     NaN
      6     5.0
      7    11.0
      8    18.0
      9    21.0
Name: clicks, dtype: float64