Aplicando função com vários argumentos para criar uma nova coluna pandas

165

Eu quero criar uma nova coluna em um pandasquadro de dados aplicando uma função a duas colunas existentes. Após esta resposta, consegui criar uma nova coluna quando preciso apenas de uma coluna como argumento:

import pandas as pd
df = pd.DataFrame({"A": [10,20,30], "B": [20, 30, 10]})

def fx(x):
    return x * x

print(df)
df['newcolumn'] = df.A.apply(fx)
print(df)

No entanto, não consigo descobrir como fazer a mesma coisa quando a função exige vários argumentos. Por exemplo, como eu crio uma nova coluna passando a coluna A e a coluna B para a função abaixo?

def fxy(x, y):
    return x * y

python pandas

— Michael
fonte

136

Como alternativa, você pode usar a função subjacente numpy:

>>> import numpy as np
>>> df = pd.DataFrame({"A": [10,20,30], "B": [20, 30, 10]})
>>> df['new_column'] = np.multiply(df['A'], df['B'])
>>> df
    A   B  new_column
0  10  20         200
1  20  30         600
2  30  10         300

ou vetorizar função arbitrária no caso geral:

>>> def fx(x, y):
...     return x*y
...
>>> df['new_column'] = np.vectorize(fx)(df['A'], df['B'])
>>> df
    A   B  new_column
0  10  20         200
1  20  30         600
2  30  10         300

— alko
fonte

2

Obrigado pela resposta! Estou curioso, esta é a solução mais rápida?

— MV23

6

A versão vetorizada np.vectorize()é incrivelmente rápida. Obrigado.

— stackoverflowuser2010

Esta é uma solução útil. Se o tamanho dos argumentos de entrada para a função xey não forem iguais, você receberá um erro. Nesse caso, a solução @RomanPekar funciona sem nenhum problema. Não comparei o desempenho.

— Ehsan Sadr

Eu sei que essa é uma resposta antiga, mas: eu tenho um caso extremo, no qual np.vectorizenão funciona. O motivo é que uma das colunas é do tipo pandas._libs.tslibs.timestamps.Timestamp, que é transformada no tipo numpy.datetime64pela vetorização. Os dois tipos não são intercambiáveis, fazendo com que a função se comporte mal. Alguma sugestão sobre isso? (Fora .applyisso, aparentemente, isso deve ser evitado)

— ElRudi

Ótima solução! caso alguém esteja se perguntando, o vetorize funciona bem e super rápido para as funções de comparação de cadeias.

— infiniteloop

227

Você pode usar o exemplo @greenAfrican, se for possível reescrever sua função. Mas se você não deseja reescrever sua função, pode envolvê-la em uma função anônima dentro de apply, assim:

>>> def fxy(x, y):
...     return x * y

>>> df['newcolumn'] = df.apply(lambda x: fxy(x['A'], x['B']), axis=1)
>>> df
    A   B  newcolumn
0  10  20        200
1  20  30        600
2  30  10        300

— Roman Pekar
fonte

4

Essa é uma ótima dica e deixa as referências da coluna próximas à chamada de aplicação (na verdade, na mesma). Usei essa dica e a dica de saída de várias colunas @toto_tico fornecida para gerar uma função de 3 colunas e 4 colunas! Funciona bem!

— RufusVS

7

Uau, parece que você é o único que não se concentra no exemplo mínimo do OP, mas resolve todo o problema, obrigado, exatamente o que eu precisava! :) #

— 1916 Matt

38

Isso resolve o problema:

df['newcolumn'] = df.A * df.B

Você também pode fazer:

def fab(row):
  return row['A'] * row['B']

df['newcolumn'] = df.apply(fab, axis=1)

— verde-africano
fonte

10

Esta resposta resolve este exemplo de brinquedo e será suficiente para eu reescrever minha função real, mas não aborda como aplicar uma função definida anteriormente sem reescrevê-la nas colunas de referência.

— 11743 Michael

23

Se você precisar criar várias colunas ao mesmo tempo :

Crie o quadro de dados:

import pandas as pd
df = pd.DataFrame({"A": [10,20,30], "B": [20, 30, 10]})

Crie a função:

def fab(row):                                                  
    return row['A'] * row['B'], row['A'] + row['B']

Atribua as novas colunas:

df['newcolumn'], df['newcolumn2'] = zip(*df.apply(fab, axis=1))

— toto_tico
fonte

1

Eu queria saber como eu poderia gerar várias colunas com uma aplicação! Usei isso com a resposta do @Roman Pekar para gerar uma função de 3 colunas e 4 colunas! Funciona bem!

— RufusVS 11/09

15

Mais uma sintaxe limpa de estilo de ditado:

df["new_column"] = df.apply(lambda x: x["A"] * x["B"], axis = 1)

ou,

df["new_column"] = df["A"] * df["B"]

— Surya
fonte