将函数应用于数据框中的列并返回两个新列的值Python

Question

I have the following DF and f(a,b) function:

       A        B

0      5       3
1      4       2
2      7       1

f(a,b):
 return (a+b,a-b)

I want to a apply f(a,b) on columns A,B ... and return two values into two new columns df[sum,sub]

       A      B       C       D

0      5      3       8       2
1      4      2       6       2
2      7      1       8       6

Answer 1

Using apply with axis=1

import pandas as pd
df = pd.DataFrame({"A": [5, 4, 7], "B":[3, 2, 1]})

def f(a,b):
    return (a+b,a-b)

df[["sum", "sub"]] = df.apply(lambda row: f(row["A"], row["B"]), axis=1).apply(pd.Series)
print(df)

Output:

   A  B  sum  sub
0  5  3    8    2
1  4  2    6    2
2  7  1    8    6

Answer 2

This is one way. I strongly recommend you don't use pd.DataFrame.apply with a row-wise calculation, as this unnecessarily sidesteps pandas vectorisation.

def f(a, b):
    return a + b, a - b

def foo(df, a, b):
    return f(df[a], df[b])

df['C'], df['D'] = df.pipe(foo, 'A', 'B')

print(df)

   A  B  C  D
0  5  3  8  2
1  4  2  6  2
2  7  1  8  6