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 的总和为 "next 1" a，c2 的总和为 "next 2"，c3 的总和为"next 3" 个 a 和 c4 的总和为 "next 4" 个 a。

因此，我想要的输出是：

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

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

Pandas: Take rolling sum of next (1 ... n) rows of a column within a group and create a new column for each sum

python

dataframe

cumulative-sum

pandas

pandas-groupby