加速 pandas 滚动 window

Question

我想加快我使用pandas.rolling().apply()自定义 function 的代码。 下面的代码运行良好，但速度非常慢。 应用数百万行时，有什么方法可以加快速度。

for i in [12, 9, 6, 3]:
    df[f'want_col_{i}'] = df.groupby(['account'])['types'].rolling(window = i).apply(lambda x: sum(x == 1)).values

这个想法是在给定的滚动 window 中计算值。 例如，从上面的代码中，我喜欢通过给定的 window 12, 9, 6, 3分别按account计算值等于1组。

有没有办法提高速度，谢谢！

Answer 1

你可以试试：

df['types_eq_1'] = df['types'].eq(1).astype(int)

for i in [12, 9, 6, 3]:
    df[f'want_col_{i}'] = df.groupby(['account'])['types_eq_1'].rolling(window = i).sum()

df = df.drop('types_eq_1', 1)