Search code examples
pythonpandasdataframerolling-computation

Pandas rolling mean with variable window based on an different column

I would like to perform the rolling mean on a window that varies depending on the values of a column in my DataFrame. Can anyone help me? Here is a starting point:

import pandas as pd
import numpy as np

rng = np.random.default_rng()
df = pd.DataFrame(rng.integers(0, 100, size=(100, 2)), columns=list('AB'))
df.loc[:,'B']=df['B']//10

Now I would like to get the rolling mean of the series df.A with the window based on column B. For example if df.B[0] is worth 3 my_series[0]=df.A.rolling(3).mean()[0] and so on for my_series[1] etc...

Can you help me? ty for your time I appreciate it.

Solution

  • One option is to loop through the data frame, and assign a new column equal to the rolling_mean for each row.

    df['rolling_mean'] = np.nan
for ind in range(len(df)):
    df.loc[df.index[ind], 'rolling_mean'] = df.A.rolling(df.loc[df.index[ind], 'B']).mean()[ind]