使用 groupBy 在 Pandas 中使用 shift 和滚动

Question

df = pd.DataFrame(dict(
    list(
        zip(["A", "B", "C"],
            [np.array(["id %02d" % i for i in range(1, 11)]).repeat(10),
            pd.date_range("2018-01-01", periods=100).strftime("%Y-%m-%d"),
            [i for i in range(10, 110)]])
        )
))

df = df.groupby(["A", "B"]).sum()

df["D"] = df["C"].shift(1).rolling(2).mean()

df

此代码生成以下内容：

我希望滚动逻辑为每个新 ID 重新开始。现在，ID 02正在使用来自的最后两个值ID 01来计算平均值。

如何做到这一点？

Answer 1

我相信你需要groupby：

df['D'] = df["C"].shift(1).groupby(df['A'], group_keys=False).rolling(2).mean()
print (df.head(20))
                   C     D
A     B                   
id 01 2018-01-01  10   NaN
      2018-01-02  11   NaN
      2018-01-03  12  10.5
      2018-01-04  13  11.5
      2018-01-05  14  12.5
      2018-01-06  15  13.5
      2018-01-07  16  14.5
      2018-01-08  17  15.5
      2018-01-09  18  16.5
      2018-01-10  19  17.5
id 02 2018-01-11  20   NaN
      2018-01-12  21  19.5
      2018-01-13  22  20.5
      2018-01-14  23  21.5
      2018-01-15  24  22.5
      2018-01-16  25  23.5
      2018-01-17  26  24.5
      2018-01-18  27  25.5
      2018-01-19  28  26.5
      2018-01-20  29  27.5

或者：

df['D'] = df["C"].groupby(df['A']).shift(1).rolling(2).mean()
print (df.head(20))
                   C     D
A     B                   
id 01 2018-01-01  10   NaN
      2018-01-02  11   NaN
      2018-01-03  12  10.5
      2018-01-04  13  11.5
      2018-01-05  14  12.5
      2018-01-06  15  13.5
      2018-01-07  16  14.5
      2018-01-08  17  15.5
      2018-01-09  18  16.5
      2018-01-10  19  17.5
id 02 2018-01-11  20   NaN
      2018-01-12  21   NaN
      2018-01-13  22  20.5
      2018-01-14  23  21.5
      2018-01-15  24  22.5
      2018-01-16  25  23.5
      2018-01-17  26  24.5
      2018-01-18  27  25.5
      2018-01-19  28  26.5
      2018-01-20  29  27.5

Answer 2

虽然 @jezrael 接受的答案对于正转变可以正确工作，但对于负转变却给出了错误的结果（部分）。请检查以下内容

df['D'] = df["C"].groupby(df['A']).shift(1).rolling(2).mean()
df['E'] = df["C"].groupby(df['A']).rolling(2).mean().shift(1).values
df['F'] = df["C"].groupby(df['A']).shift(-1).rolling(2).mean()
df['G'] = df["C"].groupby(df['A']).rolling(2).mean().shift(-1).values
df.set_index(['A', 'B'], inplace=True)
print(df.head(20))

                   C     D     E     F     G
A     B                                     
id 01 2018-01-01  10   NaN   NaN   NaN  10.5
      2018-01-02  11   NaN   NaN  11.5  11.5
      2018-01-03  12  10.5  10.5  12.5  12.5
      2018-01-04  13  11.5  11.5  13.5  13.5
      2018-01-05  14  12.5  12.5  14.5  14.5
      2018-01-06  15  13.5  13.5  15.5  15.5
      2018-01-07  16  14.5  14.5  16.5  16.5
      2018-01-08  17  15.5  15.5  17.5  17.5
      2018-01-09  18  16.5  16.5  18.5  18.5
      2018-01-10  19  17.5  17.5   NaN   NaN
id 02 2018-01-11  20   NaN  18.5   NaN  20.5
      2018-01-12  21   NaN   NaN  21.5  21.5
      2018-01-13  22  20.5  20.5  22.5  22.5
      2018-01-14  23  21.5  21.5  23.5  23.5
      2018-01-15  24  22.5  22.5  24.5  24.5
      2018-01-16  25  23.5  23.5  25.5  25.5
      2018-01-17  26  24.5  24.5  26.5  26.5
      2018-01-18  27  25.5  25.5  27.5  27.5
      2018-01-19  28  26.5  26.5  28.5  28.5
      2018-01-20  29  27.5  27.5   NaN   NaN

请注意，列D和E是针对 计算的.shift(1)，而列F和G是针对 计算的.shift(-1)。列E不正确，因为 的第一个值id 02使用 的最后两个值id 01。列F不正确，因为和 的第一个值都是NaNs 。列并给出正确的结果。所以，完整的答案应该是这样的。如果轮班周期非负，则使用以下id 01id 02DG

df['D'] = df["C"].groupby(df['A']).shift(1).rolling(2).mean()

如果轮班周期为负数，则使用以下公式

df['G'] = df["C"].groupby(df['A']).rolling(2).mean().shift(-1).values

希望能帮助到你！