Speed up pandas rolling window

Question

I want to speedup my code that I used pandas.rolling().apply() for custom function. The code below is worked fine but it is super slow. Is there any way to speedup it when applying with million of rows.

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

The idea is to count value in given rolling window. For example, from the code above I like to count value is equal to 1 group by account by given window 12, 9, 6, 3 respectively.

Is there anyway to increse the speed, thanks!

Answer 1

You could try:

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)