有条件的 pandas 中的移动平均线

Question

我有一个具有以下结构的 dataframe：

import numpy as np
import pandas as pd

df = pd.DataFrame(
    {
        "date": ["2020-01-01", "2020-01-02", "2020-01-03", "2020-01-04"] * 2,
        "group": ["A", "A", "A", "A", "B", "B", "B", "B"],
        "x": [1, 2, 2, 3, 2, 3, 4, 2],
        "condition": [1, 0, 1, 0] * 2
    }
)
df

我想计算列 x 的最后 3 天的滚动移动平均值：

• 每组
• 仅使用过去的数据（不使用当前行）
• 仅使用condition = 1的滚动平均值的数据。

结果应如下所示：

我怎样才能在 pandas 中做到这一点？ 谢谢！

请记住，它与此不同：

在 pandas 中滚动 function 有条件

在这里，我正在寻找过去 3 天的移动平均线，而在另一天，我只想要移动平均线。

Answer 1

首先用Series.where用NaN替换不匹配的行，然后按组移动值并调用滚动方法：

f = lambda x: x.shift().rolling(3, min_periods=1).mean()
df['roll'] = (df.assign(x = df['x'].where(df['condition'].eq(1)))
                .groupby('group')['x']
                .transform(f))
print (df)
         date group  x  condition  roll
0  2020-01-01     A  1          1   NaN
1  2020-01-02     A  2          0   1.0
2  2020-01-03     A  2          1   1.0
3  2020-01-04     A  3          0   1.5
4  2020-01-01     B  2          1   NaN
5  2020-01-02     B  3          0   2.0
6  2020-01-03     B  4          1   2.0
7  2020-01-04     B  2          0   3.0

详情：

print (df.assign(x = df['x'].where(df['condition'].eq(1))))
         date group    x  condition
0  2020-01-01     A  1.0          1
1  2020-01-02     A  NaN          0
2  2020-01-03     A  2.0          1
3  2020-01-04     A  NaN          0
4  2020-01-01     B  2.0          1
5  2020-01-02     B  NaN          0
6  2020-01-03     B  4.0          1
7  2020-01-04     B  NaN          0