如何用Python基于列表公式批量处理多层索引DataFrame生成新列？

给定如下Pandas DataFrame与公式列表：

import pandas as pd
cols = pd.MultiIndex.from_tuples([("K", "A1", "A"), 
                                  ("K", "B1", "B"), 
                                  ("K", "C1", "C"),
                                  ("M", "A1", "A"), 
                                  ("M", "B1", "B"), 
                                  ("M", "C1", "C")])

data = [[5, 10, 40, 4, 8, 9], [2, 15, 70, 1, 7, 3], [6, 14, 60, 12, 41, 61]]
df = pd.DataFrame(data, columns=cols)
lst = [5, '*', 'A1.A', '+', 'C1.C']

需求为解析lst中的公式（对应逻辑：5 × df['A1','A'] + df['C1','C']），并对所有第一级索引（如K、M）对应的匹配第二、三级索引的列组合执行运算，生成如下结果列：

df[("K", "N1", "N")] = 5 * df[("K", "A1", "A")] + df[("K", "C1", "C")]
df[("M", "N1", "N")] = 5 * df[("M", "A1", "A")] + df[("M", "C1", "C")]

请问该需求是否可行？

1. 解析公式列表，转换为可复用的计算组件 #

先将lst中类似A1.A的列名格式拆分，转换为MultiIndex的第二、三级索引元组，再保留运算符和数值，构建计算模板。

2. 遍历第一级索引，批量生成计算列 #

利用MultiIndex的层级特性提取所有唯一的第一级分组，为每个分组动态生成对应的计算表达式并执行，最终添加新列。

完整实现代码：

import pandas as pd

# 初始化数据
cols = pd.MultiIndex.from_tuples([("K", "A1", "A"), 
                                  ("K", "B1", "B"), 
                                  ("K", "C1", "C"),
                                  ("M", "A1", "A"), 
                                  ("M", "B1", "B"), 
                                  ("M", "C1", "C")])

data = [[5, 10, 40, 4, 8, 9], [2, 15, 70, 1, 7, 3], [6, 14, 60, 12, 41, 61]]
df = pd.DataFrame(data, columns=cols)
lst = [5, '*', 'A1.A', '+', 'C1.C']

# 解析列名字符串为(第二级, 第三级)元组
def parse_col(col_str):
    sec, third = col_str.split('.')
    return (sec, third)

# 转换公式列表为可替换的组件
formula_components = []
for item in lst:
    if isinstance(item, str) and '.' in item:
        formula_components.append(parse_col(item))
    else:
        formula_components.append(str(item))

# 获取所有唯一的第一级索引
first_level_groups = df.columns.get_level_values(0).unique()

# 为每个分组生成计算列
for group in first_level_groups:
    # 构建当前分组的计算表达式
    expr_parts = []
    for comp in formula_components:
        if isinstance(comp, tuple):
            # 替换为当前分组对应的列引用
            expr_parts.append(f"df[('{group}', '{comp[0]}', '{comp[1]}')]")
        else:
            expr_parts.append(comp)
    # 拼接并执行表达式
    calculation = eval(' '.join(expr_parts))
    # 赋值给新列
    df[(group, "N1", "N")] = calculation

# 查看结果
print(df)

代码说明 #

• parse_col：将A1.A格式的列名转换为MultiIndex层级元组，方便定位DataFrame列。
• 动态表达式构建：针对每个第一级分组，将公式组件中的列引用替换为对应分组的列，生成可执行的计算表达式。
• 批量生成列：遍历所有分组，执行计算并添加新的MultiIndex列(group, "N1", "N")，完全匹配需求。

执行后，DataFrame会新增目标列，计算结果与预期一致。

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