如何在Python pandas中按两变量分组并基于另一变量生成新列

我能在R里实现这个功能，但完全搞不懂Python里怎么弄。我有个数据集，包含sbj、num_item、visit和height字段，想用Pandas创建baseline_height列。需求是：按sbj和num_item分组，每组里的baseline_height列都设置为该组中visit为Baseline对应的height值。

示例数据 #

我尝试过的错误代码 #

下面几种写法都没成功：

df['baseline_height'] = df.groupby(by = ['sbj', 'num_item']).height[['visit' == "Baseline"]]

df['baseline_height'] = 0
df = df.loc[df.groupby(by = ['sbj', 'num_item'])['baseline_height'].apply('visit' == 'Baseline') == df['height']]

df['baseline_height'] = df.groupby(by = ['sbj', 'num_item']).apply(['height'][['visit'] == 'Baseline'])

df['baseline_height'] = df.groupby(by = ['sbj', 'num_item']).apply(df['height'][df["visit"]=='Baseline'])

df_grouped = df.groupby(by = ['sbj', 'num_item'])
df['baseline_height'] = df_grouped.height[df_grouped["visit"]=='Baseline']

正确解决方案 #

方法1：使用groupby.transform（最简洁） #

transform可以将分组计算的结果自动广播到原数据的每一行，完美匹配需求：

df['baseline_height'] = df.groupby(['sbj', 'num_item']).transform(
    lambda group: group['height'][group['visit'] == 'Baseline'].iloc[0]
)['height']

方法2：提取基线数据后合并 #

先筛选出所有基线行，再通过merge将基线值匹配到对应分组的所有行：

# 提取每组的基线身高，重命名列
baseline_data = df[df['visit'] == 'Baseline'][['sbj', 'num_item', 'height']]
baseline_data.rename(columns={'height': 'baseline_height'}, inplace=True)

# 合并到原数据集
df = df.merge(baseline_data, on=['sbj', 'num_item'], how='left')

错误原因说明 #

你之前的代码问题在于：

• 分组后直接索引height[['visit' == "Baseline"]]是语法错误，无法正确筛选分组内的行
• apply的用法不正确，没有返回每个分组对应的单个基线值，也没有处理广播到所有行的逻辑
• 尝试通过loc筛选行的思路偏离了需求，我们需要的是新增列而非过滤行

内容的提问来源于stack exchange，提问作者NICE8x