如何將多個相互關聯的列傳遞給groupby和agg上的函數？

Question

我有以下熊貓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需要幾個相互關聯的列作為輸入。 最終目標是為每個id創建一個新列new_col 。

我該怎么做？

這是我當前的代碼：

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(...??))

Answer 1

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