Application de la fonction avec plusieurs arguments pour créer une nouvelle colonne pandas

165

Je souhaite créer une nouvelle colonne dans un pandasbloc de données en appliquant une fonction à deux colonnes existantes. Suite à cette réponse, j'ai pu créer une nouvelle colonne lorsque je n'ai besoin que d'une seule colonne comme argument:

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)

Cependant, je ne peux pas comprendre comment faire la même chose lorsque la fonction nécessite plusieurs arguments. Par exemple, comment créer une nouvelle colonne en passant la colonne A et la colonne B à la fonction ci-dessous?

def fxy(x, y):
    return x * y
python pandas Michael
la source

Réponses:

136

Vous pouvez également utiliser la fonction sous-jacente 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 vectoriser une fonction arbitraire dans le cas général:

>>> 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
la source
2
Merci d'avoir répondu! Je suis curieux, est-ce la solution la plus rapide?
MV23
6
La version vectorisée utilisant np.vectorize()est incroyablement rapide. Je vous remercie.
stackoverflowuser2010
C'est une solution utile. Si la taille des arguments d'entrée de la fonction x et y n'est pas égale, vous obtenez une erreur. Dans ce cas, la solution @RomanPekar fonctionne sans aucun problème. Je n'ai pas comparé les performances.
Ehsan Sadr
Je sais que c'est une vieille réponse, mais: j'ai un cas de pointe, dans lequel np.vectorizene fonctionne pas. La raison est que l'une des colonnes est du type pandas._libs.tslibs.timestamps.Timestamp, qui est transformé en type numpy.datetime64par la vectorisation. Les deux types ne sont pas interchangeables, ce qui entraîne un mauvais comportement de la fonction. Des suggestions à ce sujet? (Autre que .applycela est apparemment à éviter)
ElRudi
Excellente solution! au cas où quelqu'un se demanderait, vectorize fonctionne bien et très rapidement pour les fonctions de comparaison de chaînes.
infiniteloop
227

Vous pouvez aller avec l'exemple @greenAfrican, s'il vous est possible de réécrire votre fonction. Mais si vous ne voulez pas réécrire votre fonction, vous pouvez l'envelopper dans une fonction anonyme à l'intérieur de apply, comme ceci:

>>> 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
la source
4
C'est un bon conseil, et cela laisse les références de colonne près de l'appel apply (en fait). J'ai utilisé cette astuce et la pointe de sortie multi-colonnes @toto_tico fournie pour générer une fonction 3 colonnes en entrée, 4 colonnes en sortie! Fonctionne très bien!
RufusVS
7
Wow, il semble que vous soyez le seul à ne pas vous concentrer sur l'exemple minimal d'OP mais à résoudre tout le problème, merci, exactement ce dont j'avais besoin! :)
Matt
38

Cela résout le problème:

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

Vous pouvez également faire:

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

df['newcolumn'] = df.apply(fab, axis=1)
vertafricain
la source
10
Cette réponse résout cet exemple de jouet et me suffira pour réécrire ma fonction réelle, mais elle ne traite pas de la façon d'appliquer une fonction précédemment définie sans la réécrire dans des colonnes de référence.
Michael
23

Si vous devez créer plusieurs colonnes à la fois :

  1. Créez le dataframe:

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

  2. Créez la fonction:

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

  3. Attribuez les nouvelles colonnes:

    df['newcolumn'], df['newcolumn2'] = zip(*df.apply(fab, axis=1))
toto_tico
la source
1
Je me demandais comment je pourrais générer plusieurs colonnes avec une seule application! J'ai utilisé cela avec la réponse de @Roman Pekar pour générer une fonction 3 colonnes en entrée, 4 colonnes en sortie! Fonctionne très bien!
RufusVS
15

Une autre syntaxe propre de style dict:

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

ou,

df["new_column"] = df["A"] * df["B"]
Surya
la source