python - 是否有更清晰的方法来应用需要多个数据框列的函数来分组数据？

Question

假设我有一个函数“f”，需要三个参数:一个数据帧和两列。

如果我想对分组数据使用“func”来创建新列。这是一个设置:

 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')

有没有更干净的方法来做到这一点？我真的希望有一种方法可以使用 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 )

你会得到相同的结果