Aplicación de funciones con múltiples argumentos para crear una nueva columna de pandas

165

Quiero crear una nueva columna en un pandasmarco de datos aplicando una función a dos columnas existentes. Después de esta respuesta , he podido crear una nueva columna cuando solo necesito una columna 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)

Sin embargo, no puedo entender cómo hacer lo mismo cuando la función requiere múltiples argumentos. Por ejemplo, ¿cómo creo una nueva columna pasando la columna A y la columna B a la función a continuación?

def fxy(x, y):
    return x * y
python pandas Miguel
fuente

Respuestas:

136

Alternativamente, puede usar la función subyacente 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

o vectorizar funciones arbitrarias en caso general:

>>> 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
fuente
2
¡Gracias por la respuesta! Tengo curiosidad, ¿es esta la solución más rápida?
MV23
66
La versión vectorizada que usa np.vectorize()es asombrosamente rápida. Gracias.
stackoverflowuser2010
Esta es una solución útil. Si el tamaño de los argumentos de entrada para la función xey no es igual, obtendrá un error. En ese caso, la solución @RomanPekar funciona sin ningún problema. No comparé el rendimiento.
Ehsan Sadr
Sé que esta es una respuesta anterior, pero: tengo un caso límite, en el que np.vectorizeno funciona. La razón es que una de las columnas es del tipo pandas._libs.tslibs.timestamps.Timestamp, que se convierte en el tipo numpy.datetime64por la vectorización. Los dos tipos no son intercambiables, lo que hace que la función se comporte mal. ¿Alguna sugerencia sobre esto? (Aparte de que .applyaparentemente esto debe evitarse)
ElRudi
Gran solución! en caso de que alguien se pregunte, vectorize también funciona bien y súper rápido para las funciones de comparación de cadenas.
infiniteloop
227

Puede ir con el ejemplo de @greenAfrican, si le es posible reescribir su función. Pero si no desea reescribir su función, puede envolverla en una función anónima dentro de aplicar, como esta:

>>> 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
fuente
44
Este es un gran consejo, y deja las referencias de columna cerca de la llamada de aplicación (en realidad). Utilicé esta sugerencia y la sugerencia de salida de varias columnas @toto_tico suministrada para generar una función de 3 columnas, 4 columnas. ¡Funciona genial!
RufusVS
77
Wow, parece que eres el único que no se enfoca en el ejemplo mínimo de OP, sino que soluciona todo el problema, ¡gracias, exactamente lo que necesitaba! :)
Matt
38

Esto resuelve el problema:

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

También puedes hacer:

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

df['newcolumn'] = df.apply(fab, axis=1)
verdeafricano
fuente
10
Esta respuesta resuelve este ejemplo de juguete y será suficiente para que reescriba mi función real, pero no aborda cómo aplicar una función definida previamente sin reescribirla en columnas de referencia.
Michael
23

Si necesita crear varias columnas a la vez :

  1. Crea el marco de datos:

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

  2. Crea la función:

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

  3. Asigne las nuevas columnas:

    df['newcolumn'], df['newcolumn2'] = zip(*df.apply(fab, axis=1))
toto_tico
fuente
1
¡Me preguntaba cómo podría generar múltiples columnas con una sola aplicación! ¡Utilicé esto con la respuesta de @Roman Pekar para generar una función de 3 columnas en 4 columnas! ¡Funciona genial!
RufusVS
15

Una sintaxis limpia de estilo dict más:

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

o,

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