创建一个前N行的新列作为数组

Question

I have a dataframe df that looks like this,

   a       b 
0  30.05  29.55
1  30.20  26.05
2  30.81  25.65
3  31.12  26.44
.. ...    ...
85 30.84  25.65
86 31.12  26.44
87 29.55  25.57
88 32.41  25.45
89 21.55  29.57
90 32.91  26.41
91 34.12  25.69

I need to create a new Column 'c' that holds an Array of Column 'b' value plus the previous 4 rows values of Column 'b'. So the resulting df would look like,

     a      b     c
0  30.05  29.55 [29.55,0,0,0,0]
1  30.20  26.05 [26.05,29.55,0,0,0]
2  30.81  25.65 [25.65,26.05,29.55,0,0]
3  31.12  26.44 [26.44,25.65,26.05,29.55,0]
.. ...    ...
85 30.84  25.65 [25.65, 44.60, 30.15, 29.55, 24.66 ]
86 31.12  26.44 [26.44, 25.65, 25.65, 25.65, 25.65 ]
87 29.55  25.57 [25.57, 26.44, 25.65, 25.65, 25.65 ]
88 32.41  25.45 [25.45, 25.57, 26.44, 25.65, 25.65 ]
89 21.55  29.57 [29.57, 25.45, 25.57, 26.44, 25.65 ]
90 32.91  26.41 [26.41, 29.57, 25.45, 25.57, 26.44 ]
91 34.12  25.69 [25.69, 26.41, 29.57, 25.45, 25.57 ]

I know I can access previous rows with df.b.shift(1) and df.b.shift(2) etc but I want to be able to easily change how many rows I look back to form the array with a variable rather than type out the many shift(n)

After looking all day I'm stuck. (python3.6)

Answer 1

You could use pd.concat with range(N)

In [60]: df['c'] = pd.concat([df.b.shift(i) for i in range(4)], 1).fillna(0).values.tolist()

In [61]: df
Out[61]:
        a      b                             c
0   30.05  29.55        [29.55, 0.0, 0.0, 0.0]
1   30.20  26.05      [26.05, 29.55, 0.0, 0.0]
2   30.81  25.65    [25.65, 26.05, 29.55, 0.0]
3   31.12  26.44  [26.44, 25.65, 26.05, 29.55]
85  30.84  25.65  [25.65, 26.44, 25.65, 26.05]
86  31.12  26.44  [26.44, 25.65, 26.44, 25.65]
87  29.55  25.57  [25.57, 26.44, 25.65, 26.44]
88  32.41  25.45  [25.45, 25.57, 26.44, 25.65]
89  21.55  29.57  [29.57, 25.45, 25.57, 26.44]
90  32.91  26.41  [26.41, 29.57, 25.45, 25.57]
91  34.12  25.69  [25.69, 26.41, 29.57, 25.45]

Or , use np.column_stack on shift(n)

In [70]: np.column_stack([df.b.shift(i).fillna(0) for i in range(4)]).tolist()
Out[70]:
[[29.55, 0.0, 0.0, 0.0],
 [26.05, 29.55, 0.0, 0.0],
 [25.65, 26.05, 29.55, 0.0],
 [26.44, 25.65, 26.05, 29.55],
 [25.65, 26.44, 25.65, 26.05],
 [26.44, 25.65, 26.44, 25.65],
 [25.57, 26.44, 25.65, 26.44],
 [25.45, 25.57, 26.44, 25.65],
 [29.57, 25.45, 25.57, 26.44],
 [26.41, 29.57, 25.45, 25.57],
 [25.69, 26.41, 29.57, 25.45]]

Answer 2

You can use a conditional list comprehension (to check when the lookback is before the first value in the index).

rows_lookback = 5

df = df.assign(c=[[df['b'].iat[n - i] if n - i >= 0 else 0 
                   for i in range(rows_lookback)] 
                  for n in range(len(df['b']))])
>>> df
        a      b                                    c
0   30.05  29.55                  [29.55, 0, 0, 0, 0]
1   30.20  26.05              [26.05, 29.55, 0, 0, 0]
2   30.81  25.65          [25.65, 26.05, 29.55, 0, 0]
3   31.12  26.44      [26.44, 25.65, 26.05, 29.55, 0]
85  30.84  25.65  [25.65, 26.44, 25.65, 26.05, 29.55]
86  31.12  26.44  [26.44, 25.65, 26.44, 25.65, 26.05]
87  29.55  25.57  [25.57, 26.44, 25.65, 26.44, 25.65]
88  32.41  25.45  [25.45, 25.57, 26.44, 25.65, 26.44]
89  21.55  29.57  [29.57, 25.45, 25.57, 26.44, 25.65]
90  32.91  26.41  [26.41, 29.57, 25.45, 25.57, 26.44]
91  34.12  25.69  [25.69, 26.41, 29.57, 25.45, 25.57]