Pandas rolling window to return an array

Here is a sample code.

df = pd.DataFrame(np.random.randn(10, 2), columns=list('AB'))
df['C'] = df.B.rolling(window=3)

Output:

           A         B                                       C
0 -0.108897  1.877987  Rolling [window=3,center=False,axis=0]
1 -1.276055 -0.424382  Rolling [window=3,center=False,axis=0]
2  1.578561 -1.094649  Rolling [window=3,center=False,axis=0]
3 -0.443294  1.683261  Rolling [window=3,center=False,axis=0]
4  0.674124  0.281077  Rolling [window=3,center=False,axis=0]
5  0.587773  0.697557  Rolling [window=3,center=False,axis=0]
6 -0.258038 -1.230902  Rolling [window=3,center=False,axis=0]
7 -0.443269  0.647107  Rolling [window=3,center=False,axis=0]
8  0.347187  0.753585  Rolling [window=3,center=False,axis=0]
9 -0.369179  0.975155  Rolling [window=3,center=False,axis=0]

I want my ‘C’ column to be an array like [0.1231, -1.132, 0.8766]. I tried using rolling apply but in vain.

Expected Output:

       A         B                 C
0 -0.108897  1.877987  []
1 -1.276055 -0.424382  []
2  1.578561 -1.094649  [-1.094649, -0.424382, 1.877987]
3 -0.443294  1.683261  [1.683261, -1.094649, -0.424382]
4  0.674124  0.281077  [0.281077, 1.683261, -1.094649]
5  0.587773  0.697557  [0.697557, 0.281077, 1.683261]
6 -0.258038 -1.230902  [-1.230902, 0.697557, 0.281077]
7 -0.443269  0.647107  [0.647107, -1.230902, 0.697557]
8  0.347187  0.753585  [0.753585, 0.647107, -1.230902]
9 -0.369179  0.975155  [0.975155, 0.753585, 0.647107]

Answer

You could use np.stride_tricks:

import numpy as np
as_strided = np.lib.stride_tricks.as_strided  

df

          A         B
0 -0.272824 -1.606357
1 -0.350643  0.000510
2  0.247222  1.627117
3 -1.601180  0.550903
4  0.803039 -1.231291
5 -0.536713 -0.313384
6 -0.840931 -0.675352
7 -0.930186 -0.189356
8  0.151349  0.522533
9 -0.046146  0.507406

win = 3  # window size

# https://stackoverflow.com/a/47483615/4909087
v = as_strided(df.B, (len(df) - (win - 1), win), (df.B.values.strides * 2))

v
array([[ -1.60635669e+00,   5.10129842e-04,   1.62711678e+00],
       [  5.10129842e-04,   1.62711678e+00,   5.50902812e-01],
       [  1.62711678e+00,   5.50902812e-01,  -1.23129111e+00],
       [  5.50902812e-01,  -1.23129111e+00,  -3.13383794e-01],
       [ -1.23129111e+00,  -3.13383794e-01,  -6.75352179e-01],
       [ -3.13383794e-01,  -6.75352179e-01,  -1.89356194e-01],
       [ -6.75352179e-01,  -1.89356194e-01,   5.22532550e-01],
       [ -1.89356194e-01,   5.22532550e-01,   5.07405549e-01]])

df['C'] = pd.Series(v.tolist(), index=df.index[win - 1:])
df

          A         B                                                  C
0 -0.272824 -1.606357                                                NaN
1 -0.350643  0.000510                                                NaN
2  0.247222  1.627117  [-1.606356691642917, 0.0005101298424200881, 1....
3 -1.601180  0.550903  [0.0005101298424200881, 1.6271167809032248, 0....
4  0.803039 -1.231291  [1.6271167809032248, 0.5509028122535129, -1.23...
5 -0.536713 -0.313384  [0.5509028122535129, -1.2312911105674484, -0.3...
6 -0.840931 -0.675352  [-1.2312911105674484, -0.3133837943758246, -0....
7 -0.930186 -0.189356  [-0.3133837943758246, -0.6753521794378446, -0....
8  0.151349  0.522533  [-0.6753521794378446, -0.18935619377656243, 0....
9 -0.046146  0.507406  [-0.18935619377656243, 0.52253255045267, 0.507...