How to create lists from pandas columns

Question

I have created a pandas dataframe using this code:

import numpy as np
import pandas as pd

ds = {'col1': [1,2,3,3,3,6,7,8,9,10]}


df = pd.DataFrame(data=ds)

The dataframe looks like this:

print(df)

   col1
0     1
1     2
2     3
3     3
4     3
5     6
6     7
7     8
8     9
9    10

I need to create a field called col2 that contains in a list (for each record) the last 3 elements of col1 while iterating through each record. So, the resulting dataframe would look like this:

Does anyone know how to do it by any chance?

Answer 1

Here is a solution using rolling and list comprehension

df['col2'] = [x.tolist() for x in df['col1'].rolling(3)]

   col1        col2
0     1         [1]
1     2      [1, 2]
2     3   [1, 2, 3]
3     3   [2, 3, 3]
4     3   [3, 3, 3]
5     6   [3, 3, 6]
6     7   [3, 6, 7]
7     8   [6, 7, 8]
8     9   [7, 8, 9]
9    10  [8, 9, 10]