在Pandas多层列DataFrame中对0级列执行自定义计算

问题背景 #

现有一个带多级列索引的Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    ('A', 'a'): [1, 2, 3],
    ('A', 'b'): [4, 5, 6],
    ('B', 'a'): [7, 8, 9],
    ('B', 'b'): [10, 11, 12],
})

需要对每个0级列下的对应子列执行计算A * B / 2，生成新的0级标签为C的列，最终得到包含原列和新计算列的DataFrame。

解决方案 #

方法一：按子列分组计算（匹配你的groupby思路） #

这里需要按**子列（level=1）**分组，而非0级列，这样才能确保每个子列下的A、B列对应计算：

# 按子列分组，对每组内的A、B列执行计算
c_cols = df.groupby(level=1, axis=1).apply(lambda x: x[('A', x.name)] * x[('B', x.name)] / 2)
# 为新列设置多级索引，0级为'C'，1级保留原有的子列名
c_cols.columns = pd.MultiIndex.from_tuples([('C', col) for col in c_cols.columns])
# 合并原DataFrame与新计算列
result_df = pd.concat([df, c_cols], axis=1)

方法二：直接按子列匹配计算（更高效） #

如果子列的对应关系明确，也可以直接通过索引匹配计算，无需groupby：

# 获取所有唯一的子列名
sub_cols = df.columns.get_level_values(1).unique()
# 构造新列的字典，键为多级索引，值为计算结果
c_data = {('C', sub): df[('A', sub)] * df[('B', sub)] / 2 for sub in sub_cols}
# 将新列合并到原DataFrame
result_df = pd.concat([df, pd.DataFrame(c_data)], axis=1)

最终结果 #

运行上述任意方法后，result_df的输出如下：

A     B       C     
    a  b   a   b    a    b
0  1  4   7  10  3.5  20.0
1  2  5   8  11  8.0  27.5
2  3  6   9  12 13.5  36.0

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