Pandas按id分组取每组前2条记录计算均值并新增列的实现方法

实现代码 #

你可以通过Pandas的groupby结合transform方法完成需求，完整可运行代码如下：

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame({
    'id': [1, 1, 1, 2, 2, 2, 2, 3],
    'value': [1, 2, 3, 4, 3, 2, 1, 1]
})

# 核心计算逻辑：按id分组取每组前2条value求平均，广播到对应分组的所有行
df['top_2_mean'] = df.groupby('id')['value'].transform(lambda x: x.head(2).mean())

# 输出结果验证
print(df)

代码逻辑说明 #

• groupby('id')：根据id字段将数据划分为不同分组
• transform：区别于普通聚合方法只返回和分组数相同的结果，transform会将聚合结果按分组映射回原数据的每一行，自动保证索引对齐
• 匿名函数lambda x: x.head(2).mean()：对每个分组的value序列，取前2条记录计算平均值，若分组内记录不足2条，则自动取全部分组数据计算，刚好适配id=3的场景

运行后输出的结果和要求的格式完全一致：

id  value  top_2_mean
0   1      1         1.5
1   1      2         1.5
2   1      3         1.5
3   2      4         3.5
4   2      3         3.5
5   2      2         3.5
6   2      1         3.5
7   3      1         1.0

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