Применение функции к столбцам в dataframe и возврат к значениям двух новых столбцов Python

Question

У меня есть следующие функции DF и f (a, b):

       A        B

0      5       3
1      4       2
2      7       1

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

Я хочу применить f (a, b) к столбцам A, B ... и вернуть два значения в два новых столбца df [sum, sub]

       A      B       C       D

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

Answer 1

Это один из способов.Я настоятельно рекомендую вам не использовать pd.DataFrame.apply с построчным вычислением, так как это излишне обходит pandas векторизацию.

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

Answer 2

Использование apply с 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)

Выход:

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