Pandas 数据框 groupby 应用具有可变参数数量的函数

Question

我有一个看起来像这样的 pandas 数据框

import pandas as pd

data = {
  "Race_ID": [2,2,2,2,2,5,5,5,5,5,5],
  "Student_ID": [1,2,3,4,5,9,10,2,3,6,5],
  "theta": [8,9,2,12,4,5,30,3,2,1,50]
}

df = pd.DataFrame(data)

并且我有一个函数

f(thetai, *theta) = thetai ** 2 + the other thetas in the same race

，我想将其应用于按

theta

分组的数据框中的

Race_ID

列，并创建一个名为

feature

的新列。

所以我们有

对于第 2 场比赛中的学生 1，该值为 8^2 + 9+2+12+4

对于第 2 场比赛中的学生 2，该值为 9^2 + 8+2+12+4

对于第 2 场比赛中的学生 3，该值为 2^2 + 8+9+12+4

等等

我了解

groupby

和

apply

方法，但我不知道当参数数量可能变化时如何应用这些方法。

所以期望的结果看起来像这样

data = {
  "Race_ID": [2,2,2,2,2,5,5,5,5,5,5],
  "Student_ID": [1,2,3,4,5,9,10,2,3,6,5],
  "theta": [8,9,2,12,4,5,30,3,2,1,50],
  "fearure": [91,107,37,167,47,111,961,97,93,91,2541]
}

df = pd.DataFrame(data)

Answer 1

只需将其他 theta 的总和作为中间列即可。

theta_sum = df.groupby("Race_ID").sum()["theta"].to_dict()
# (theta sum of race_id) - self theta
df["theta_feature"] = df.apply(lambda x: theta_sum[x["Race_ID"]] - x["theta"], axis=1)
df["feature"] = df.apply(lambda x: x["theta"] ** 2 + x["theta_feature"], axis=1)

然后，如果您不想要不需要的列，您可以将其删除

df.drop(columns="theta_feature", inplace=True)