如何用自定义函数修改MultiIndex DataFrame的userid层级值？

问题场景 #

现有一个带多层列索引（MultiIndex）的DataFrame，层级包含userid、city、year等，需要仅对userid层级的索引值应用自定义函数修改，其余层级保持不变。示例DataFrame构造代码如下：

import pandas as pd
import numpy as np

arrays = [np.array(['John', 'John', 'John', 'Jane', 'Jane', 'Jane']),
          np.array(['New York', 'New York', 'San Francisco', 'New York', 'New York', 'San Francisco']),
          np.array(['2018', '2019', '2018', '2019', '2020', '2020']),
          np.array(['HR', 'Finance', 'HR', 'Finance', 'HR', 'Finance']),
          np.array(['Company A', 'Company A', 'Company A', 'Company B', 'Company B', 'Company B']),
          np.array(['Manager 1', 'Manager 1', 'Manager 2', 'Manager 2', 'Manager 3', 'Manager 3'])]

df = pd.DataFrame(np.random.randn(3, 6), columns=arrays)
df.columns.names = ['userid', 'city', 'year', 'department', 'company', 'line_manager']

自定义修改函数（注意：原示例中starts_with为错误写法，Python字符串方法应为startswith，已修正）：

def change_name(name):
    if name.startswith("J"): 
        # 仅在特定条件下修改名称
        result = "name00"
    else:
        result = name
    return result

解决方法 #

可以利用Pandas的MultiIndex.set_levels方法，直接针对指定层级修改值，同时保留原有层级的组合关系，代码简洁高效：

# 定位userid层级的位置
userid_level = 'userid'
# 对该层级的所有值应用自定义函数，并用set_levels替换原层级
df.columns = df.columns.set_levels(
    [change_name(name) for name in df.columns.levels[df.columns.names.index(userid_level)]],
    level=userid_level
)

原理说明 #

• df.columns.names.index(userid_level)：获取userid层级在MultiIndex中的索引位置
• [change_name(name) for name in ...]：对该层级的所有唯一值应用自定义函数
• set_levels(..., level=userid_level)：将修改后的值替换到指定层级，其余层级的索引值和组合关系完全保留

如果需要处理更复杂的批量修改逻辑，也可以通过提取层级数据重新构造MultiIndex的方式实现，效果一致：

# 提取原MultiIndex的层级和编码
levels = list(df.columns.levels)
codes = df.columns.codes
# 获取userid层级的位置
userid_idx = df.columns.names.index('userid')
# 修改该层级的值
levels[userid_idx] = list(map(change_name, levels[userid_idx]))
# 重新构造MultiIndex并赋值
df.columns = pd.MultiIndex(levels=levels, codes=codes, names=df.columns.names)

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