Pandas为多级索引分组DataFrame添加Total/Avg汇总行的高效方法

问题场景 #

需要在pandas多级索引的分组聚合结果中新增标识为Total/Avg的汇总行：

• 求和类列取对应上级分组维度的总和
• 均值类列取对应上级分组维度的平均值
• 要求方案适配多列不同聚合规则，避免手动重索引、逐列循环填充的繁琐操作，支持后续快速新增聚合逻辑

测试基础代码 #

测试DataFrame构造 #

import numpy as np
import pandas as pd
import random

years = [2000, 2001, 2002]
products = ["A", "B", "C"]
num_combos = 10

years = [random.choice(years) for i in range(num_combos)]
products = [random.choice(products) for i in range(num_combos)]
sum_values = list(range(0, num_combos))
random.shuffle(sum_values)
av_values = [random.randrange(0, num_combos, 1) for i in range(num_combos)]

cols = {"years": years,
        "products": products,
        "sum_col": sum_values,
        "av_col": av_values}
df = pd.DataFrame(cols)

基础分组聚合逻辑 #

按products、years两级维度分组，sum_col求和、av_col求平均：

gdf = df.groupby(["products", "years"]).agg(s = ("sum_col", "sum"),
                                            a = ("av_col", "mean"))

原有实现的缺陷 #

原有方案需要手动判断索引层级生成新索引、重命名索引层级、逐列判断聚合规则填充值，代码冗余度高，列数增加、聚合规则扩展（比如新增中位数、计数）时需要修改大量逻辑，扩展性极差。原有繁琐实现参考：

def addTotalAvgMultiindex(df):
    num_indexes = len(list(df.index.levels))
    if num_indexes == 3:
        a, b, c = df.index.levels
        df = df.reindex(pd.MultiIndex.from_product([a, b, [*c, 'Total/Avg']]))
    elif num_indexes == 4:
        a, b, c, d = df.index.levels
        df = df.reindex(pd.MultiIndex.from_product([a, b, c, [*d, 'Total/Avg']]))
    elif num_indexes == 2:
        a, b = df.index.levels
        df = df.reindex(pd.MultiIndex.from_product([a, [*b, 'Total/Avg']]))
    return df

gdf = addTotalAvgMultiindex(gdf)
gdf.index = gdf.index.set_names(["products", "years"])

for col in gdf.columns:
    if col == "s":
        total = df.groupby(["products"]).agg(total=("sum_col", "sum"))
    elif col == "a":
        total = df.groupby(["products"]).agg(total=("av_col", "mean"))
    
    total_values = [x for xs in total.values for x in xs]
        
    gdf[col][gdf.index.get_level_values("years") == "Total/Avg"] = total_values

简洁通用实现方案 #

核心思路是聚合规则一次定义、多次复用，分别计算细粒度分组结果和上级维度汇总结果，给汇总行的最内层索引打上Total/Avg标签后直接合并，完全避免手动处理索引和逐列填充的操作。

快速实现代码 #

# 统一定义聚合规则，后续新增聚合逻辑直接在该字典中添加即可
agg_rules = {
    "s": ("sum_col", "sum"),
    "a": ("av_col", "mean")
}

# 计算细粒度分组结果
gdf_detail = df.groupby(["products", "years"]).agg(**agg_rules)
# 按上级维度（products）计算汇总值，给最内层索引赋值Total/Avg
gdf_total = df.groupby("products").agg(**agg_rules)
gdf_total = gdf_total.assign(years="Total/Avg").set_index("years", append=True)
# 合并结果并按索引排序
gdf = pd.concat([gdf_detail, gdf_total]).sort_index()

通用封装函数 #

适配任意层级的多级索引场景，不需要手动判断索引层级数量：

def add_group_total(df, group_cols, total_label="Total/Avg", **agg_rules):
    """
    给多级索引分组结果追加最内层维度的汇总行
    :param df: 原始未分组DataFrame
    :param group_cols: 分组字段列表，顺序与多级索引层级顺序一致，最内层分组字段放在最后
    :param total_label: 汇总行的索引显示标签
    :param agg_rules: 聚合规则，格式与pandas agg命名聚合完全一致
    """
    # 计算全维度细粒度分组
    detail_res = df.groupby(group_cols).agg(**agg_rules)
    # 排除最内层分组字段，按上级维度计算汇总值
    upper_group_cols = group_cols[:-1]
    total_res = df.groupby(upper_group_cols).agg(**agg_rules)
    # 给汇总结果补充最内层索引标签
    total_res = total_res.assign(**{group_cols[-1]: total_label}).set_index(group_cols[-1], append=True)
    # 合并结果并排序
    return pd.concat([detail_res, total_res]).sort_index()

函数调用示例

gdf = add_group_total(
    df,
    group_cols=["products", "years"],
    s=("sum_col", "sum"),
    a=("av_col", "mean"),
    # 后续新增聚合规则直接追加参数即可，比如新增av_col的中位数计算
    # a_median = ("av_col", "median")
)

方案优势 #

• 不需要手动判断索引层级、手动构造MultiIndex做重索引操作
• 聚合规则只需要定义一次，新增列、新增聚合逻辑（中位数、计数、最大值等）时只需要补充聚合规则参数，不需要修改核心计算逻辑
• 不存在逐列循环、赋值的冗余操作，两次分组自动按照相同规则计算结果，合并后直接得到对齐的完整数据集
• 适配2层、3层及以上任意层级的多级索引分组场景，只需要调整传入的group_cols参数即可

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