如何对MultiIndex DataFrame按不同层级应用不同分组规则？

问题描述 #

现有如下多索引DataFrame：

import pandas as pd
import numpy as np

arrays = [
    ["bar", "bar", "baz", "baz", "foo", "foo", "qux", "qux"],
    ["one1", "one2", "one1", "one2", "one1", "two", "one1", "two"],
]
index = pd.MultiIndex.from_arrays(arrays, names=["first", "second"])
df = pd.DataFrame({"A": [1, 1, 1, 1, 2, 2, 3, 3], "B": np.arange(8)}, index=index)

原始DataFrame输出：

A  B
first second      
bar   one1     1  0
      one2     1  1
baz   one1     1  2
      one2     1  3
foo   one1     2  4
      two     2  5
qux   one1     3  6
      two     3  7

需求：

• 保留first（level=0）层级的原索引分组
• 对second（level=1）层级按前3个字符分组（如one1、one2合并为one），并将新索引设为该前缀
• 最终得到指定的聚合结果

解决方案 #

方法一：仅通过groupby配合聚合实现 #

可以直接在groupby中指定双层分组键，其中第二层使用索引前缀作为分组依据，配合聚合函数即可完成需求：

result = df.groupby(
    [
        df.index.get_level_values("first"),  # 按原first索引分组
        lambda x: x[:3]  # 提取second索引的前3个字符作为分组键
    ],
    names=["first", "second"]  # 设置新索引的名称
).agg({
    "A": "sum",
    "B": "sum"
})

方法二：先修改索引再聚合 #

如果觉得直接在groupby中定义分组键不够直观，也可以先重置索引并修改second列取值，再进行分组聚合：

# 重置索引为普通列
df_reset = df.reset_index()
# 修改second列为前3个字符
df_reset["second"] = df_reset["second"].str[:3]
# 分组聚合后恢复索引结构
result = df_reset.groupby(["first", "second"]).agg({"A": "sum", "B": "sum"}).sort_index()

两种方法都能得到目标结果，方法一更贴合“仅通过groupby相关语句实现”的需求。

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