如何在df.groupby之后将数据框列值作为窗口大小传递？

Question

    A   B   C
0   1   10  2
1   1   15  2
2   1   14  2
3   2   11  4
4   2   12  4
5   2   13  4
6   2   16  4
7   1   18  2

这是我的示例 DataFrame。

1. 我想在“A”列上应用 groupby，

2. 根据列 'C' 的值在列 'B' 上应用滚动总和，这意味着当 A 为 1 时，窗口大小应为 2 而不是 NaN 我想要剩余值的总和，而不管窗口大小如何。

目前我的输出是：

A   
1  0    25.0
   1    29.0
   2    32.0
   7     NaN
2  3    23.0
   4    25.0
   5    29.0
   6     NaN

上面的代码： df['B'].groupby(df['A']).rolling(df['C'][0]).sum().shift(-1)

当 C = 4 时，我希望滚动窗口为 4 并且不想要 NaN

所需的输出应如下所示：

    A   B   C   Rolling_sum
0   1   10  2   25
1   1   15  2   29
2   1   14  2   32
7   1   18  2   18
3   2   11  4   52
4   2   12  4   41
5   2   13  4   29
6   2   16  4   16

Answer 1

因为您想按列传递动态窗口，所以C使用 lambda 函数并更改顺序iloc[::-1]：

df = df.sort_values('A')
df['Rolling_sum'] = (df.iloc[::-1].groupby('A')
                       .apply(lambda x: x.B.rolling(x.C.iat[0], min_periods=0).sum())
                       .reset_index(level=0, drop=True))
print (df)
   A   B  C  Rolling_sum
0  1  10  2         25.0
1  1  15  2         29.0
2  1  14  2         32.0
7  1  18  2         18.0
3  2  11  4         52.0
4  2  12  4         41.0
5  2  13  4         29.0
6  2  16  4         16.0

如果性能很重要，则采用跨步解决方案（取决于组数、组大小、实际数据中的最佳测试）：

def rolling_window(a, window):
    a = np.concatenate([[0] * (window - 1), a])
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides).sum(axis=1)

df = df.sort_values('A')
df['Rolling_sum']  = (df.iloc[::-1].groupby('A')
                        .apply(lambda x: pd.Series(rolling_window(x.B, x.C.iat[0]), 
                                                   index=x.index))
                        .reset_index(level=0, drop=True))
print (df) 
   A   B  C  Rolling_sum
0  1  10  2           25
1  1  15  2           29
2  1  14  2           32
7  1  18  2           18
3  2  11  4           52
4  2  12  4           41
5  2  13  4           29
6  2  16  4           16