Есть ли более чистый способ применения функции, которая требует нескольких столбцов данных для сгруппированных данных?

Question

Допустим, у меня есть функция 'f', которая требует три аргумента: один фрейм данных и два столбца.

Если я хочу использовать 'fun c' для сгруппированных данных, чтобы создать новый столбец. Вот настройка:

 import pandas as pd

 def f(x,y,z):
    return sum(x[y]**2) + sum(x[z]**3) - 50

 my_df = pd.DataFrame({'state':['Alabama','Alaska','Arizona','Alaska'],'a':[1,2,3,4],'b':[4,5,6,7]})

Реальная функция, которую я пытаюсь применить, намного сложнее, но этот пример должен подойти. Я делал это:

x = my_df.groupby('state').apply(f,'a','b').rename('new_col')

my_df = my_df.merge(x, on = 'state', how = 'left')

Обычно я просто использовал бы преобразование, но преобразование выдает ошибку, потому что он может использоваться только для одного столбца за раз.

my_df.groupby('state').transform(f,'a','b') 

KeyError: ('a', 'occurred at index a')

Есть ли более чистый способ сделать это? На самом деле, sh был способ использовать apply и получить преобразованный результат, длина которого равна исходному (разгруппированному) фрейму данных.

Желаемый результат:

     state  a  b  new_col
0  Alabama  1  4       15
1   Alaska  2  5      438
2  Arizona  3  6      175
3   Alaska  4  7      438

Answer 1

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

def f(row):
    return row["a"]**2 + row["b"]**3 - 50

my_df['new_col'] = my_df.apply(f, axis=1 )

, вы получите те же результаты