pandas DataFrame中创建依赖其他变量的均值列的两种实现

原始示例数据 #

import numpy as np
import pandas as pd

df = pd.DataFrame({'Group':['a','a','b','b','b','c','c'], 'Label':[0,1,0,1,1,0,1], 'Num':[1,2,3,4,5,6,7]})

需求1：仅同组Label=1的行赋值对应均值，其余行赋值0/NaN #

# 第一步：生成每个Group对应Label=1的Num均值映射表
label1_mean_map = df[df['Label'] == 1].groupby('Group')['Num'].mean()
# 第二步：匹配映射值到所有行
df['mean'] = df['Group'].map(label1_mean_map)
# 第三步：将非Label=1的行置为0，要NaN就把0换成np.nan
df.loc[df['Label'] != 1, 'mean'] = 0

执行后df['mean']结果为[0,2,0,4.5,4.5,0,7]，符合预期。

需求2：同组所有行均赋值该组Label=1的均值 #

# 生成每个Group对应Label=1的Num均值映射表
label1_mean_map = df[df['Label'] == 1].groupby('Group')['Num'].mean()
# 直接匹配映射值到所有行即可
df['mean'] = df['Group'].map(label1_mean_map)

执行后df['mean']结果为[2,2,4.5,4.5,4.5,7,7]，符合预期。

如果偏好transform写法，也可以用下面的代码实现需求2，效果一致：

df['mean'] = df.groupby('Group').transform(lambda g: g.loc[g['Label'] == 1, 'Num'].mean())['Num']

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