如何基于指定列创建分组多级表头重塑DataFrame？

需要将DataFrame按第一列（col1）分组，为其他列添加分组表头行，实现指定的多级表头结构。

初始DataFrame #

import pandas as pd

df = pd.DataFrame(
    {
        'col1':['A','A','A','B','B','B'],
        'col2':[1,2,3,4,5,6],
        'col3':[1,2,3,4,5,6],
        'col4':[1,2,3,4,5,6],
        'colx':[1,2,3,4,5,6]
    }
)

尝试pd.pivot()的问题 #

使用pivot得到的分组层级方向不符合预期：

df.pivot(columns='col1', values=['col2','col3','col4','colx'])

输出结果：

col2      col3      col4      colx     
col1    A    B    A    B    A    B    A    B
0     1.0  NaN  1.0  NaN  1.0  NaN  1.0  NaN
1     2.0  NaN  2.0  NaN  2.0  NaN  2.0  NaN
2     3.0  NaN  3.0  NaN  3.0  NaN  3.0  NaN
3     NaN  4.0  NaN  4.0  NaN  4.0  NaN  4.0
4     NaN  5.0  NaN  5.0  NaN  5.0  NaN  5.0
5     NaN  6.0  NaN  6.0  NaN  6.0  NaN  6.0

预期输出 #

A                               B
col1    col2    col3    col4    colx    col2    col3    col4    colx
0       1       1       1       1       4       4       4       4
1       2       2       2       2       5       5       5       5
2       3       3       3       3       6       6       6       6

解决方案 #

通过给分组添加组内行索引，结合set_index和unstack调整列层级，实现预期结构：

# 为每个分组生成组内行号
df['group_idx'] = df.groupby('col1').cumcount()

# 设置多级索引并重塑，将col1转为列的上层分组
result = df.set_index(['group_idx', 'col1']).unstack('col1')

# 交换列的层级顺序，让分组（A/B）成为上层表头
result = result.swaplevel(axis=1).sort_index(axis=1)

# 移除索引名称，贴近预期输出格式
result.index.name = None

print(result)

输出结果：

A                               B
col1    col2 col3 col4 colx col2 col3 col4 colx
0         1    1    1    1    4    4    4    4
1         2    2    2    2    5    5    5    5
2         3    3    3    3    6    6    6    6

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