Применение функции к двум столбцам Pandas DataFrame

Одной из распространенных задач при работе с данными в Pandas является применение определенной функции к одному или нескольким столбцам DataFrame. Иногда возникает необходимость применить функцию, которая принимает два аргумента, к двум столбцам DataFrame, и создать новый столбец на основе полученных результатов.

В качестве примера рассмотрим следующий DataFrame:

import pandas as pd

df = pd.DataFrame({
    'A': [1, 2, 3],
    'B': [4, 5, 6],
})

Пусть есть функция, которая принимает два аргумента и возвращает их сумму:

def add_two_numbers(x, y):
    return x + y

Требуется применить эту функцию к столбцам ‘A’ и ‘B’ и сохранить результат в новом столбце ‘C’.

Самым простым и эффективным способом сделать это является использование метода apply DataFrame. Однако, при попытке применить функцию напрямую к двум столбцам, возникает ошибка о несоответствии количества аргументов:

df['C'] = df[['A', 'B']].apply(add_two_numbers)
# TypeError: ("add_two_numbers() missing 1 required positional argument: 'y'", 'occurred at index A')

Это происходит потому, что метод apply передает каждый столбец в функцию как отдельный аргумент, а функция ожидает два аргумента.

Чтобы решить эту проблему, можно использовать параметр axis метода apply и установить его значение равным 1. Это означает, что функция будет применяться к каждой строке DataFrame, а не к каждому столбцу:

df['C'] = df.apply(lambda row: add_two_numbers(row['A'], row['B']), axis=1)

Теперь новый столбец ‘C’ будет содержать сумму значений столбцов ‘A’ и ‘B’:

   A  B  C
0  1  4  5
1  2  5  7
2  3  6  9

Таким образом, при применении функций к нескольким столбцам Pandas DataFrame, следует учитывать, как функция ожидает получать свои аргументы, и соответствующим образом использовать параметр axis метода apply.