如何优雅地仅对Pandas DataFrame最后三列执行Unstack操作？

原始数据 #

import numpy as np
import pandas as pd
data = pd.DataFrame(data=np.random.rand(10,5), columns=['headA','headB','tailA','tailB','tailC'])

需求说明 #

需要仅对最后三列（tailA、tailB、tailC）执行结构重塑，生成以pd.MultiIndex.from_product([data.columns[-3:], data.index])为新索引的结果，同时保留前两列的对应值。

优雅实现方案 #

方法一：使用melt重塑数据 #

melt是pandas中宽表转长表的常用工具，可快速将指定列展开为行：

# 提取目标列并执行melt处理
result = data.melt(
    id_vars=['headA', 'headB'],  # 保留作为标识的前两列
    value_vars=data.columns[-3:],  # 需要展开的最后三列
    var_name='tail_col',  # 新索引第一层级的名称
    value_name='value'
).set_index(['tail_col', data.index.repeat(3)]).sort_index()

方法二：使用stack操作（贴合Unstack反向逻辑） #

如果想更贴近Unstack的反向思路，可通过stack实现：

# 先将前两列设为索引，再对目标列执行stack
result = (
    data.set_index(['headA', 'headB'])
    .stack()  # 将最后三列的列名转为索引层级
    .swaplevel(0, 1)  # 交换索引层级，让tail列名处于第一层级
    .sort_index()
)

验证索引匹配 #

通过以下代码可确认生成的索引与目标完全一致：

target_index = pd.MultiIndex.from_product([data.columns[-3:], data.index])
print(result.index.equals(target_index))  # 输出True

内容的提问来源于stack exchange，提问作者principal-ideal-domain