如何将多个相互关联的列传递给groupby和agg上的函数?

时间:2019-07-18 09:33:19

标签: python pandas

我有以下熊猫DataFrame df

id  col1   col2
1   7      1.2
1   6      0.8
1   12     0.9
1   1      1.1
2   3      2.0
2   6      1.8
3   10     0.7
3   11     0.9
3   12     1.2

以下是创建此df的代码:

import pandas as pd
df = pd.DataFrame({'id': [1,1,1,1,2,2,3,3,3], 
                   'col1': [7,6,12,1,3,6,10,11,12],
                   'col2': [1.2,0.8,0.9,1.1,2.0,1.8,0.7,0.9,1.2]})

我需要按id分组,并将功能myfunc应用于每个分组。问题是myfunc需要几个相互关联的列作为输入。最终目标是为每个new_col创建一个新列id

我该怎么办?

这是我当前的代码:

def myfunc(df, col1, col2):

    df1 = col1
    df2 = df[df[col2] < 1][[col1]]
    var1 = df1.iloc[0]
    var2 = df2.iloc[0][0]

    result = var2 - var1

    return result


df["new_col"] = df.groupby("id").agg(myfunc(...??))

1 个答案:

答案 0 :(得分:0)

在groupby-apply中,my_func()被传递给整个组以及所有列。您可以简单地从该组中选择列:

def myfunc(g):
    var1 = g['col1'].iloc[0]
    var2 = g.loc[g['col2'] > 1, 'col1'].iloc[0]

    return var1 / var2

df['new_col'] = df.groupby("id").apply(myfunc)