How to do pandas rolling window in both forward and backward at the same time

user_17786602 0
Viewed 112

I have a pd.DataFrame df with one column, say:

A = [1,2,3,4,5,6,7,8,2,4]
df = pd.DataFrame(A,columns = ['A'])

For each row, I want to take previous 2 values, current value and next 2 value (a window= 5) and get the sum and store it in new column. Desire output,

   A  A_sum
  1      6
  2     10
  3     15
  4     20
  5     25
  6     30
  7     28
  8     27
  2     21
  4     14

I have tried,

  1. df['A_sum'] = df['A'].rolling(2).sum()
  2. Tried with shift, but all doing either forward or backward, I'm looking for a combination of both.
pandas python-3.x python
3 Answers

Use rolling by 5, add parameter center=True and min_periods=1 to Series.rolling:

df['A_sum'] =  df['A'].rolling(5, center=True, min_periods=1).sum()
print (df)
   A  A_sum
0  1    6.0
1  2   10.0
2  3   15.0
3  4   20.0
4  5   25.0
5  6   30.0
6  7   28.0
7  8   27.0
8  2   21.0
9  4   14.0
user_29010020

If you are allowed to use numpy, then you might use numpy.convolve to get desired output

import numpy as np
import pandas as pd
A = [1,2,3,4,5,6,7,8,2,4]
B = np.convolve(A,[1,1,1,1,1], 'same')
df = pd.DataFrame({"A":A,"A_sum":B})
print(df)

output

   A  A_sum
0  1      6
1  2     10
2  3     15
3  4     20
4  5     25
5  6     30
6  7     28
7  8     27
8  2     21
9  4     14
user_107859750

You can use shift for this (straightforward if not elegant):

df["A_sum"] = df.A + df.A.shift(-2).fillna(0) + df.A.shift(-1).fillna(0) + df.A.shift(1).fillna(0)

output:

   A  A_sum
0  1    6.0
1  2   10.0
2  3   14.0
3  4   18.0
4  5   22.0
5  6   26.0
6  7   23.0
7  8   21.0
8  2   14.0
9  4    6.0
user_50329410
Related