Brian Bien Brian Bien - 15 days ago 6
Python Question

pandas dataframe sum of shift(x) for x in range(1, n)

I have a dataframe with like this, and want to add a new column that is the equivalent of applying

shift
n times. For example, let n = 2:

df = pd.DataFrame(numpy.random.randint(0, 10, (10, 2)), columns=['a','b'])

a b
0 0 3
1 7 0
2 6 6
3 6 0
4 5 0
5 0 7
6 8 0
7 8 7
8 4 4
9 2 2

df['c'] = df['b'].shift(1) + df['b'].shift(2)

a b c
0 0 3 NaN
1 7 0 NaN
2 6 6 3.0
3 6 0 6.0
4 5 0 6.0
5 0 7 0.0
6 8 0 7.0
7 8 7 7.0
8 4 4 7.0
9 2 2 11.0


In this manner, column
c
gets the sum of the previous
n
values from column
b
.

Other than a loop, is there a better way to accomplish this for a large
n
?

Answer

You can use the rolling() method with a window of 2:

df['c'] = df.b.rolling(window = 2).sum().shift()

df
    a   b     c
0   0   3   NaN
1   7   0   NaN
2   6   6   3.0
3   6   0   6.0
4   5   0   6.0
5   0   7   0.0
6   8   0   7.0
7   8   7   7.0
8   4   4   7.0
9   2   2   11.0
Comments