Início / coding / Perguntas / 79411167
Accepted
Sun Jar
Asked: 2025-02-04 17:28:01 +0800 CST

Pandas aplicam função retorna uma lista para uma nova coluna

  • 772

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.

python

2 respostas

  1. Best Answer

    Você pode usar numpy's sliding_window_view:

    from numpy.lib.stride_tricks import sliding_window_view as swv

N = 3
df['C'] = pd.Series(swv(df['B'], N).tolist(), index=df.index[N-1:])

    Saída:

       A  B          C
0  4  9        NaN
1  0  2        NaN
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]
    • 6

  2. Acho que você pode mudar seu pensamento de outra forma, digamos, não, row-wisemas com column-wisejanelas deslizantes, e provavelmente seu código poderia ficar mais rápido, a menos que você tenha um tamanho de janela grande N.

    Por exemplo, você pode tentar

    N = 3
nr = len(df)
df['C'] = [None]*(N-1) + np.column_stack([df['B'].iloc[k:nr-N+1+k] for k in range(N)]).tolist()

    e você obterá

        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]
    • 2

relate perguntas