pandas dataframe 带有 groupby 的滑动窗口

Question

我可以添加一个新的列c，它是b中最后两个值的总和，如下所示...

df['c'] = df.b.rolling(window = 2).sum().shift()

df
    a   b     c
0   1   3   NaN
1   1   0   NaN
2   0   6   3.0
3   1   0   6.0
4   0   0   6.0
5   1   7   0.0
6   0   0   7.0
7   0   7   7.0
8   1   4   7.0
9   1   2   11.0

...然而，如果我想首先按a分组怎么办？例如，我可以这样做：

df['c'] = df.groupby(['a'])['b'].shift(1) + df.groupby(['a'])['b'].shift(2)

有没有更优雅的方法来对一组进行大量的位移求和（1、2、...n）？

- Brian Bien

Answer 1

f = lambda x: x.rolling(2).sum().shift()
df['c'] = df.groupby('a').b.apply(f)

df

- piRSquared

---

注意：在lambda函数中结合rolling()和shift()方法（就像piRSquared所展示的那样）是必要的：它会导致两者都应用于组（这是期望的）;在这种情况下会出现错误行为：df['c'] = df.groupby('a').b.rolling(2).sum().shift()，因为shift()操作发生在非分组上下文中。 - Brian Bien

抱歉，我希望我没有增加困惑：我的意思是说你的方法是正确的，而另一种方法，可能看起来是语法上的偏好，会导致意外的行为。 - Brian Bien