如何将DataFrame中各X分组的首个元素替换为对应均值？

给定数据 #

我们有如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame(
    {'X': ['a', 'a', 'b', 'a', 'b'],
     'Y': [2, 4, 8, 10, 5]})

显示结果：

X   Y
0   a   2
1   a   4
2   b   8
3   a   10
4   b   5

需求 #

按X列分组，将每个分组的首个元素的Y值替换为该分组的均值，预期输出：

X   Y
0   a   5.33
1   a   4.00
2   b   6.50
3   a   10.00
4   b   5.00

解决方案 #

完整实现代码 #

import pandas as pd

df = pd.DataFrame(
    {'X': ['a', 'a', 'b', 'a', 'b'],
     'Y': [2, 4, 8, 10, 5]})

# 1. 计算每个分组的Y列均值
group_means = df.groupby('X')['Y'].mean()

# 2. 给每个分组内的行添加组内序号（从0开始）
df['group_rank'] = df.groupby('X').cumcount()

# 3. 替换每组首个元素（group_rank=0）的Y值为对应分组均值
df.loc[df['group_rank'] == 0, 'Y'] = df.loc[df['group_rank'] == 0, 'X'].map(group_means)

# 4. 格式化Y列为两位小数，匹配预期输出格式
df['Y'] = df['Y'].round(2)

# 移除临时辅助列
df = df.drop('group_rank', axis=1)

print(df)

代码说明 #

• groupby('X')['Y'].mean()：按X分组计算每组Y的平均值，得到一个以X取值为索引的Series，存储对应分组的均值。
• groupby('X').cumcount()：为每个分组内的行生成从0开始的连续序号，每组的第一行序号为0，这是我们要定位的目标行。
• df.loc[df['group_rank'] == 0, 'Y'] = ...：通过布尔索引定位所有分组的首个元素，用map方法根据X的值匹配对应的分组均值，完成替换。
• round(2)：将Y列数值保留两位小数，让输出格式和预期一致。

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