如何向滚动窗口添加条件？

Question

我有以下熊猫数据框df ：

COL1   COL2   COL3   Y
10     2      A      1
20     5      A      3
30     2      B      1
20     7      B      4
15     2      A      2
25     1      B      1
10     3      A      1
25     1      A      1

我将滚动应用于y如下：

window = 2
y = df["Y"]
y = y.rolling(window).apply(lambda x: np.max(x) if len(x)>0 else 0).dropna()

但是现在我需要对y添加一个限制： max应该只在COL3等于A行上计算。 如果行中没有A值，则y应该等于0。例如第3行和第4行（如果我们使用2的window ）

我试过：

y = df.rolling(window).apply(lambda row: np.max(row[row["COL3"=="A"]]["Y"]) if len(row["Y"])>0 else 0).dropna()["Y"]

但是得到了错误：

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

Answer 1

我们可以在rolling之前拆分 y 并用 0 reindex填充值

y1 = y[df.COL3 == 'A']
y1 = y1.rolling(window).apply(lambda x: np.max(x) if len(x)>0 else 0).fillna('drop')
y = y1.reindex(y.index, fill_value = 0).loc[lambda x : x!='drop']