在pandas中使用shift和rolling与groupBy一起使用

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的最后两个值来计算平均值。

如何实现这一点？

- Nilzone-

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

- jezrael

---

我现在测试了没有 group_keys=False 参数，它仍然可以工作。我真的需要这个吗？ - Nilzone-

我加上它是因为在第一种解决方案中得到了NaN，但如果正常工作，则可以自由省略它。 - jezrael

@Nilzone - 在第二种解决方案上进行了测试，如果省略它对我很有效。 - jezrael

是的 - 至少在这种情况下看起来它不会影响结果 :) - Nilzone-

@jezrael，我发现你的解决方案在负移位周期方面不起作用（至少部分）。在原始问题中，周期为正，但为了全面了解情况，我认为应该添加负周期的正确选项。请查看我的答案。 - mskoryk

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

请注意，对于.shift(1)，列D和E是计算出来的，而对于.shift(-1)，列F和G是计算出来的。由于id 02的第一个值使用了id 01的最后两个值，因此列E是不正确的。由于id 01和id 02的第一个值都是NaN，所以列F是不正确的。列D和G给出了正确的结果。因此，完整的答案应该像这样。如果移位周期为非负数，则应使用以下内容。

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

如果shift period是负数，则使用以下内容

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

希望能够帮助到你！

- mskoryk

---

这很有帮助！很容易忽略这个差异！ - Zed Fang