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

Time:04-01

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.

CodePudding user response:

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

CodePudding user response:

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

CodePudding user response:

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
  • Related