python - Pandas：如何在 groupby 对象上使用自定义应用函数返回多列

Question

基本思想是我的计算涉及数据框中的多个列并返回多个列，我想将它们集成到数据框中。我想做这样的事情：

df = pd.DataFrame({'id':['i1', 'i1', 'i2', 'i2'], 'a':[1,2,3,4], 'b':[5,6,7,8]})

def custom_f(a, b):
    computation = a+b
    return computation + 1, computation*2

df['c1'], df['c2'] = df.groupby('id').apply(lambda x: custom_f(x.a, x.b))

期望的输出：

    id  a   b  c1     c2
0   i1  1   5  7      12
1   i1  2   6  9      16
2   i2  3   7  11     20
3   i2  4   8  13     24

我知道如何一次完成这一列，但实际上使用两列的“计算”操作非常昂贵，所以我试图弄清楚我如何只能运行一次。

编辑：我意识到给定的示例可以在没有 groupby 的情况下解决，但是对于我的实际“计算”用例，我正在使用 groupby，因为我使用每个组中数组的第一个和最后一个值进行计算. 为了简单起见，我省略了它，但想象一下它是必需的。

标签： pythonpandaspandas-groupbyapply

Answer 1

df['c1'], df['c2'] = custom_f(df['a'], df['b']) # you dont need apply for your desired output here