Pandas 在组内滚动总和

Question

我正在尝试计算每组内的滚动总和或任何其他统计数据（例如平均值）。下面我给出一个例子，窗口为2，统计量为sum。

df = pd.DataFrame.from_dict({'class': ['a', 'b', 'b', 'c', 'c', 'c', 'b', 'a', 'b'],
                        'val': [1, 2, 3, 4, 5, 6, 7, 8, 9]})
df['sum2_per_class'] = [1, 2, 5, 4, 9, 11, 10, 9, 16] # I want to compute this column
# df['sum2_per_class'] = df[['class', 'val']].groupby('class').rolling(2).sum() # what I tried

    class  val  sum2_per_class
 0     a    1               1
 1     b    2               2
 2     b    3               5
 3     c    4               4
 4     c    5               9
 5     c    6              11
 6     b    7              10
 7     a    8               9
 8     b    9              16

这是我尝试过的以及相应的错误：

df['sum2_per_class'] = df[['class', 'val']].groupby('class').rolling(2).sum()

TypeError: incompatible index of inserted column with frame index

Answer 1

正如错误消息所示，滚动求和操作返回一个带有 MultiIndex 的 pandas Series，它不能直接分配给数据帧中的单个列。

一个可能的修复方法是使用

reset_index()

将 MultiIndex 转换为普通索引，如下所示：

df['sum2_per_class'] = df[['class', 'val']].groupby('class').rolling(2).sum().reset_index(level=0, drop=True)

但是，在运行上述代码后，我在“sum2_per_class”列中得到了意外的 NaN 值，如下所示：

[NaN, NaN, 5, NaN, 9, 11, 10, 9, 16]

，而其他值则符合预期。

调查 NaN 问题后，我得出以下结论：

• 滚动求和运算需要每组中至少有两个CONSECUTIVE行来计算总和。例如，对于第一组 'a'，我们有： 1)

  Row 0

  与

  val1=1

  和

  Row 7

  与

  val=8

  您期望滚动总和为

  1 + 8 = 9

  ，而这些行不连续并导致 NaN。对于我们获得预期滚动总和的其他组，分组的行是连续的。例如，对于“c”组，我们有：

  Row 3

  、

  Row 4

  和

  Row 5

  。