如何在DataFrame中按组计算均值并分配Loop ID以生成多边形

原始数据 #

import pandas as pd

mydict = {"year" : [2010, 2010, 2011, 2011, 2010, 2010, 2011, 2011],
          "region" : [1,1,1,1,2,2,2,2],
          "group" : ["lower", "upper", "lower", "upper", "lower", "upper", "lower", "upper"],
          "var" : [10,20,30,40,50,60,70,80]}
df = pd.DataFrame(mydict)

需求 #

• 新增average列：按年份计算var的均值，每条记录对应其所属年份的均值
• 新增loop列：按区域分配环形配对编号，规则为：先按年份升序排列lower组，编号从0开始递增；再按年份降序排列upper组，编号从区域数据总条数减1开始递减，形成环形结构

目标效果 #

newdict = {"year" : [2010, 2010, 2011, 2011, 2010, 2010, 2011, 2011],
          "region" : [1,1,1,1,2,2,2,2],
          "group" : ["lower", "upper", "lower", "upper", "lower", "upper", "lower", "upper"],
          "var" : [10,20,30,40,50,60,70,80],
          "average" : [15,15,35,35,55,55,75,75],
          "loop" : [0,3,1,2,0,3,1,2]}
pd.DataFrame(newdict)

实现代码 #

# 生成average列
df['average'] = df['year'].map(df.groupby('year')['var'].mean())

# 生成loop列
def process_region(group):
    # 按年份升序、组名升序排序，确保lower组在前
    sorted_df = group.sort_values(['year', 'group'], ascending=[True, True])
    total = len(sorted_df)
    # 前半部分lower组编号从0开始，后半部分upper组从末尾倒序
    sorted_df['loop'] = list(range(total//2)) + list(range(total-1, total//2 -1, -1))
    # 恢复原数据的索引顺序
    return sorted_df.reindex(group.index)

df = df.groupby('region').apply(process_region).reset_index(drop=True)

print(df)

代码说明 #

• average列：先通过groupby计算每个年份的var均值，再用map将均值匹配到对应年份的每一行
• loop列：按区域分组后，先对组内数据排序，保证lower组按年份升序排列，然后给lower组分配递增编号；upper组则从区域数据总条数的最后一位开始倒序编号，最后恢复原数据的顺序，确保编号对应正确

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