CODEWITHSUNDEEP
pythonpandasrolling-average
emax
emax

Reputation: 7235

Python: how to properly use rolling mean in a loop with Pandas?

I have a dataframe df containing the information of the transaction between 2 companies and the time. I have to group-by every 3 months and make a comparison with other 3 months plus 1. For instance I have to group together October 2015, November 2015, December 2015 and compare them with November 2015, December 2015, January 2016. So, I have to group together the months [201510, 201511, 201512] and compare them with [201511, 201512, 201601]

The dataframe looks like the following:

     A          B            YM       tot
0   494     6.83353e+07     201507  136388.22
1   1150    6.78366e+07     201507  68972.76
2   1575    6.96231e+07     201507  43447.37
3   3459    1.70194e+07     201507  298173.15
4   8591    5.40416e+07     201507  51255.22
5   17350   1.79459e+07     201507  24400.00
6   24685   1.7862e+07      201507  67631.19
7   28157   1.79105e+07     201507  20241.00
8   47963   2.73774e+07     201507  30000.00

times = pd.unique(df['YM']) ##  months we consider

times:  
array([201507, 201508, 201509, 201510, 201511, 201512, 201601, 201602,
           201603, 201604, 201605, 201606, 201607, 201608, 201609, 201610,
           201611, 201612, 201701, 201702, 201703, 201704, 201705, 201706,
           201707, 201708, 201709, 201710, 201711, 201712])

this what I am doing:

k = 0
v = 3
for i in range(0, len(times)-3)
    ## First Time Window
    tmp = df[(df['YM'] >= times[k]) & (df['YM'] <= times[v])] 
    net1 = net1.groupby(['A','B'], as_index = False)['tot'].sum()

    ## Second Time Window
    tmp = df[(df['YM'] >= times[k+1]) & (df['YM'] <= times[v+1])]  
    net2 = net2.groupby(['A','B'], as_index = False)['tot'].sum()

    k += 1 ## Update Time windows
    v += 1

I would like to know if there is a more efficient way to do that.

Upvotes: 0

Views: 1205

Answers (1)

sacuL
sacuL

Reputation: 51335

In your sample data, we only have one YM to play with, so it doesn't look like much, but I think this might do what you're looking for:

df['YM'] = pd.to_datetime(df['YM'], format='%Y%m')

df.groupby('YM').sum().rolling(freq='M', window=3).mean()

It groups by year and month, gets the sum, and then gets a rolling mean of each 3 months

If you want to limit the comparison to the tot column:

df.groupby('YEARMONTH')['tot'].sum().rolling(freq='M', window=3).mean()

Upvotes: 2

Related Questions