Add n values from the both sides of a pandas data frame column values

user_11629296 0
Viewed 139

I have a data frame like this,

df
col1    col2
  A       1
  B       3
  C       2
  D       5
  E       6
  F       8
  G       10

I want to add previous and next n values of a particular value of col2 and store it into a new column,

So, If n=2, then the data frame should look like,

 col1    col2    col3
  A       1       6  (only below 2 values are there no upper values, so adding 3 numbers)
  B       3      11 (adding one prev, current and next two)
  C       2      17(adding all 4 values)
  D       5      24(same as above)
  E       6      31(same as above)
  F       8      29(adding two prev and next one as only one is present)
  G       10     24(adding with only prev two values)

When previous or next 2 values are not found adding whatever values are available. I can do it using a for loop, but the execution time will be huge, looking for some pandas shortcuts do do it most efficiently.

rolling-sum pandas dataframe python
1 Answers

You can use the rolling method.

import pandas as pd
df = pd.read_json('{"col1":{"0":"A","1":"B","2":"C","3":"D","4":"E","5":"F","6":"G"},"col2":{"0":1,"1":3,"2":2,"3":5,"4":6,"5":8,"6":10}}')

df['col3'] = df['col2'].rolling(5, center=True, min_periods=0).sum()

col1    col2    col3
0   A   1   6.0
1   B   3   11.0
2   C   2   17.0
3   D   5   24.0
4   E   6   31.0
5   F   8   29.0
6   G   10  24.0
user_76485780
Related