Pandas MultiIndex指定规则排序与分组统计问题求解

问题场景 #

给定双层MultiIndex结构的DataFrame，构造代码如下：

import numpy as np
import pandas as pd

arrays = [
    np.array(["baz", "baz", "bar", "bar", "qux",  "foo"]),
    np.array(["yes", "no", "yes", "no", "yes",   "no"]),
]
df = pd.DataFrame(np.random.randint(100, size=(6,4)), index=arrays)

需求为按一级索引分组，计算每列的yes_rate（即yes行数值/分组内yes、no行数值总和），同时满足两个要求：

• 索引按固定规则排序：一级索引顺序保持原始出现顺序[baz,bar,qux,foo]，二级索引顺序固定为[no,yes,All,yes_rate]
• 代码重复执行时统计值不漂移，计算总和时仅统计yes、no原始行，排除已生成的All、yes_rate统计行，同时兼容分组内yes/no行不同时存在的场景

修复后实现代码 #

# 预定义索引排序规则
FIRST_LEVEL_ORDER = ["baz", "bar", "qux", "foo"]
SECOND_LEVEL_ORDER = ["no", "yes", "All", "yes_rate"]

for first_idx in FIRST_LEVEL_ORDER:
    # 仅筛选当前分组下二级索引为yes/no的有效行计算总和，自动过滤历史生成的统计行
    group_data = df.loc[first_idx]
    valid_data = group_data[group_data.index.isin(["yes", "no"])]
    group_sum = valid_data.sum()

    # 存在yes行时计算yes_rate
    if "yes" in valid_data.index:
        df.loc[(first_idx, "yes_rate"), :] = valid_data.loc["yes"] / group_sum
    
    # 写入分组汇总值
    df.loc[(first_idx, "All"), :] = group_sum

# 将两级索引转为有序分类类型，实现自定义排序
df.index = pd.MultiIndex.from_arrays([
    pd.CategoricalIndex(df.index.get_level_values(0), categories=FIRST_LEVEL_ORDER, ordered=True),
    pd.CategoricalIndex(df.index.get_level_values(1), categories=SECOND_LEVEL_ORDER, ordered=True)
])
df = df.sort_index()

核心逻辑说明 #

• 解决重复执行统计错误问题：求和前先过滤出二级索引为yes/no的行，无论之前是否生成过统计行，都只会用原始业务数据计算，同时自动适配yes/no行缺失的场景，不需要额外逐行判断存在性再累加
• 解决自定义排序问题：将索引转为带固定顺序的分类类型后，sort_index()会严格按照指定优先级排序，不会默认按字典序排列，完全匹配要求的索引顺序

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