Eu tenho um dataframe pandas:
import pandas as pd
import numpy as np
np.random.seed(150)
df = pd.DataFrame(np.random.randint(0, 10, size=(10, 2)), columns=['A', 'B'])
Quero adicionar uma nova coluna "C" cujos valores são a lista combinada de cada três linhas na coluna "B". Então, uso o método a seguir para atingir minhas necessidades, mas esse método é lento quando os dados são grandes.
>>> df['C'] = [df['B'].iloc[i-2:i+1].tolist() if i >= 2 else None for i in range(len(df))]
>>> df
A B C
0 4 9 None
1 0 2 None
2 4 5 [9, 2, 5]
3 7 9 [2, 5, 9]
4 8 3 [5, 9, 3]
5 8 1 [9, 3, 1]
6 1 4 [3, 1, 4]
7 4 1 [1, 4, 1]
8 1 9 [4, 1, 9]
9 3 7 [1, 9, 7]
Quando tento usar a função df.apply, recebo uma mensagem de erro:
df['C'] = df['B'].rolling(window=3).apply(lambda x: list(x), raw=False)
TypeError: must be real number, not list
Lembro que o pandas apply não parece retornar uma lista, então há uma maneira melhor? Procurei no fórum, mas não consegui encontrar um tópico adequado sobre apply e return.
