Pandas：为多层索引列的DataFrame添加新计算列

问题背景 #

现有一个带多级列索引的DataFrame，结构及初始化代码如下：

import numpy as np
import pandas as pd

np.random.seed(1)
mi = pd.MultiIndex.from_product([[3, 5], ["X","Y","V","T"]], names=["Node", "Parameter"])
df = pd.DataFrame(index=pd.DatetimeIndex(['2022-07-07 12:00:00', '2022-07-07 13:00:00',
               '2022-07-07 14:00:00', '2022-07-07 15:00:00',
               '2022-07-07 16:00:00'],
             dtype='datetime64[ns]', name='Date', freq=None), columns=mi, data=np.random.rand(5,8))

输出的DataFrame结构：

Node                        3                                       5
Parameter                   X         Y         V         T         X         Y         V         T
Date
2022-07-07 12:00:00  0.417022  0.720324  0.000114  0.302333  0.146756  0.092339  0.186260  0.345561
2022-07-07 13:00:00  0.396767  0.538817  0.419195  0.685220  0.204452  0.878117  0.027388  0.670468
2022-07-07 14:00:00  0.417305  0.558690  0.140387  0.198101  0.800745  0.968262  0.313424  0.692323
2022-07-07 15:00:00  0.876389  0.894607  0.085044  0.039055  0.169830  0.878143  0.098347  0.421108
2022-07-07 16:00:00  0.957890  0.533165  0.691877  0.315516  0.686501  0.834626  0.018288  0.750144

需求是为每个Node（3和5）添加新计算列Z，计算规则为X² + Y²。已有实现方式如下：

x = df.loc[:,(slice(None),"X")]
y = df.loc[:,(slice(None),"Y")]
z = (x**2).rename(columns={"X":"Z"}) + (y ** 2).rename(columns={"Y":"Z"})
result = df.join(z).sort_index(axis=1)

希望找到更简洁的实现方式，同时了解如何用df.xs选取数据后赋值到原DataFrame。

简洁实现方法 #

方法1：按Node分组批量计算 #

直接按第一级列索引（Node）分组，自动处理所有Node的X/Y计算，代码更紧凑：

# 按Node分组，每组内计算X²+Y²
z_cols = df.groupby(level="Node", axis=1).apply(lambda g: g["X"]**2 + g["Y"]**2)
# 重构多级列索引，对应(Node, Z)格式
z_cols.columns = pd.MultiIndex.from_tuples(
    [(node, "Z") for node in z_cols.columns], 
    names=["Node", "Parameter"]
)
# 合并原数据并排序列索引
result = pd.concat([df, z_cols], axis=1).sort_index(axis=1)

这种方法无需手动拆分X/Y列，新增Node时无需修改代码，扩展性更强。

方法2：用df.xs选取数据后直接赋值 #

xs可以快速提取指定层级的列，计算后直接赋值到新的多级列索引位置：

# 复制原DataFrame避免修改原始数据
result = df.copy()
# 遍历所有唯一Node
for node in df.columns.get_level_values("Node").unique():
    # 提取当前Node的X、Y列
    x = df.xs((node, "X"), axis=1)
    y = df.xs((node, "Y"), axis=1)
    # 计算Z并赋值到对应的多级列位置
    result[(node, "Z")] = x**2 + y**2
# 按列索引排序
result = result.sort_index(axis=1)

逻辑直观，适合需要对单个Node做额外处理的场景。

方法3：用eval简化计算逻辑 #

利用df.eval简化表达式，结合扁平列名过渡后再转回多级索引：

# 将多级列索引转为扁平格式（如"3_X"）
flat_df = df.copy()
flat_df.columns = [f"{n}_{p}" for n,p in flat_df.columns]
# 批量计算每个Node的Z值
for node in [3,5]:
    flat_df.eval(f"{node}_Z = {node}_X**2 + {node}_Y**2", inplace=True)
# 将扁平列名转回多级索引
flat_df.columns = pd.MultiIndex.from_tuples(
    [(int(c.split("_")[0]), c.split("_")[1]) for c in flat_df.columns],
    names=["Node", "Parameter"]
)
result = flat_df.sort_index(axis=1)

代码简洁，适合计算规则简单的场景。

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