Pandas：获取组内列的下 (1 ... n) 行的滚动总和，并为每个总和创建一个新列

Question

我有以下数据框：

a = [1,2,3,4,5,6,7,8]
x1 = ['j','j','j','k','k','k','k','k']
df = pd.DataFrame({'a': a,'b':x1})  

print(df)

a   b
1   j
2   j
3   j
4   k
5   k
6   k
7   k
8   k

我正在尝试获取“b”列中分组的下 n 行的“a”值的总和，并将其存储在新列中（n 的范围从 1 到 4）。

基本上我希望得到四个新列 c1、c2、c3 和 c4，这样 c1 的总和为“下一个 1”a，c2 的总和为“下一个 2”a，c3 的总和为“下一个 3” " a's 和 c4 的总和为 "next 4" a's。

因此，我想要的输出是：

a   b   c1      c2      c3      c4  
1   j   2.0     5.0     NaN     NaN
2   j   3.0     NaN     NaN     NaN
3   j   NaN     NaN     NaN     NaN
4   k   5.0     11.0    18.0    26.0
5   k   6.0     13.0    21.0    NaN
6   k   7.0     15.0    NaN     NaN
7   k   8.0     NaN     NaN     NaN
8   k   NaN     NaN     NaN     NaN

我寻找解决方案，我能想到的最好的方法是：

for x in range(1,5): 
    df[x] = df.groupby(['b'])a[::-1].rolling(x+1).sum()[::-1] - a

但是这种语法会引发错误。

如果可能的话，如果我需要按多个字段进行分组，您能否也分享一下如何实现。非常感谢任何帮助。

谢谢。

Answer 1

您的示例数据框与您的预期输出不匹配，所以让我们使用后者。

我认为您可以将滚动总和与移位相结合：

for x in range(1, 5):
    c = pd.Series(df.groupby("b")["a"].rolling(x).sum().values, index=df.index)
    df[f"c{x}"]= c.groupby(df["b"]).shift(-x)

给我

In [302]: df
Out[302]: 
   a  b   c1    c2    c3    c4
0  1  j  2.0   5.0   NaN   NaN
1  2  j  3.0   NaN   NaN   NaN
2  3  j  NaN   NaN   NaN   NaN
3  4  k  5.0  11.0  18.0  26.0
4  5  k  6.0  13.0  21.0   NaN
5  6  k  7.0  15.0   NaN   NaN
6  7  k  8.0   NaN   NaN   NaN
7  8  k  NaN   NaN   NaN   NaN

如果你真的想拥有多个键，你可以使用一个键列表，但我们必须稍微重新安排调用：

keys = ["b","b2"]
for x in range(1, 5):
    c = pd.Series(df.groupby(keys)["a"].rolling(x).sum().values, index=df.index)
    df[f"c{x}"]= c.groupby([df[k] for k in keys]).shift(-x)

或

keys = ["b","b2"]
for x in range(1, 5):
    c = pd.Series(df.groupby(keys)["a"].rolling(x).sum().values, index=df.index)
    df[f"c{x}"]= df.assign(tmp=c).groupby(keys)["tmp"].shift(-x)

给我

In [409]: df
Out[409]: 
   a  b b2   c1    c2  c3  c4
0  1  j  j  2.0   5.0 NaN NaN
1  2  j  j  3.0   NaN NaN NaN
2  3  j  j  NaN   NaN NaN NaN
3  4  k  k  5.0   NaN NaN NaN
4  5  k  k  NaN   NaN NaN NaN
5  6  k  l  7.0  15.0 NaN NaN
6  7  k  l  8.0   NaN NaN NaN
7  8  k  l  NaN   NaN NaN NaN