如何筛选Pandas多级索引行，实现索引层级间的1:1对应关系

需求：对多级索引的DataFrame进行筛选，每个一级索引仅保留一行，筛选逻辑为保留该行所有列值求和最大的行。

初始数据与目标效果 #

以下是生成初始DataFrame（start_df）和目标DataFrame（end_df）的代码：

import numpy as np
import pandas as pd

# 初始DataFrame
tuples = [
    (1, 1),
    (2, 9), (2, 4), (2, 3),
    (3, 2), (3, 11)
]

start_df = pd.DataFrame(
    {
       'col1':  [1,1,2,1,2,1],
       'col2': [1,1,1,1,2,1]
    },
    index=pd.MultiIndex.from_tuples(tuples)
)

# 目标DataFrame
tuples = [
    (1, 1),
    (2, 4), 
    (3, 2), 
]

end_df = pd.DataFrame(
    {
       'col1':  [1,2,2],
       'col2': [1,1,2]
    },
    index=pd.MultiIndex.from_tuples(tuples)
)

实现方案 #

通过以下步骤实现需求：

1. 新增临时列存储每行的列值总和
2. 按一级索引分组，筛选出每组中总和最大的行
3. 删除临时列，得到最终结果

完整代码：

import numpy as np
import pandas as pd

# 构建初始数据
tuples = [
    (1, 1),
    (2, 9), (2, 4), (2, 3),
    (3, 2), (3, 11)
]

start_df = pd.DataFrame(
    {
       'col1':  [1,1,2,1,2,1],
       'col2': [1,1,1,1,2,1]
    },
    index=pd.MultiIndex.from_tuples(tuples)
)

# 1. 计算每行的列值总和
start_df['row_sum'] = start_df.sum(axis=1)

# 2. 按一级索引分组，保留每组中总和最大的行
result_df = start_df.groupby(level=0).apply(lambda x: x[x['row_sum'] == x['row_sum'].max()]).droplevel(0)

# 3. 删除临时总和列
result_df = result_df.drop('row_sum', axis=1)

# 验证结果与目标一致
print(result_df.equals(end_df))  # 输出: True

关键说明 #

• sum(axis=1)：按行计算所有列值的总和，axis=1指定行维度求和
• groupby(level=0)：基于多级索引的第一级进行分组
• droplevel(0)：移除分组后新增的外层索引，恢复原多级索引结构

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