有條件的 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