如何将分组列的长格式pd.DataFrame转换为宽格式/方阵格式

我有一个DataFrame，每行最多包含4组观测数据，每组数据是(id, x, y)的组合，id是唯一标识，x和y是对应数值。现在需要将数据重塑，让每个id对应两列（x_{id}和y_{id}），每行中没有对应id观测的位置用NA填充。试过各种pivot组合都没成功，以下是示例数据：

import pandas as pd

df_before = pd.DataFrame(
    [
        [5, 'a', '0', 9, 'a', '1'],
        [5, 'b', '3', 6, 'c', '3'],
        [9, 'c', '4'],
        [2, 'd', '12', 6, 'a', '3', 9, 'c', '2', 4, 'a', '3'],
        [6, 'a', '3'],
        [5, 'b', '0', 9, 'a', '4', 4, 'd', '9'],
        [2, 'c', '2', 9, 'g', '5', 4, 'g', '8']
    ],
    columns=['id_0','x_0','y_0', 'id_1', 'x_1', 'y_1', 'id_2', 'x_2', 'y_2', 'id_3', 'x_3', 'y_3']
).fillna(pd.NA)

df_after = pd.DataFrame(
    [
        [pd.NA, pd.NA, pd.NA, pd.NA, 'a', '0', pd.NA, pd.NA, 'a', '1'],
        [pd.NA, pd.NA, pd.NA, pd.NA, 'b', '3', 'c', '3', pd.NA, pd.NA],
        [pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, 'c', '4'],
        ['d', '12', 'a', '3', pd.NA, pd.NA, 'a', '3', 'c', '2'],
        [pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, pd.NA, 'a', '3', pd.NA, pd.NA],
        [pd.NA, pd.NA, 'd', '9', 'b', '0', pd.NA, pd.NA, 'a', '4'],
        ['c', '2', 'g', '8', pd.NA, pd.NA, pd.NA, pd.NA, 'g', '5']
    ],
    columns=['x_2', 'y_2', 'x_4', 'y_4', 'x_5', 'y_5', 'x_6', 'y_6', 'x_9', 'y_9']
)

解决方案 #

通过先转长格式、构造目标列名、再转回宽格式的流程可以实现需求，具体代码如下：

import pandas as pd

# 1. 将原宽表转为长格式，保留原行索引
df_long = df_before.unstack().reset_index()
df_long.columns = ['col', 'idx', 'value']

# 拆分列名，提取分组序号和数据类型（id/x/y）
df_long['group'] = df_long['col'].str.split('_').str[1]
df_long['type'] = df_long['col'].str.split('_').str[0]

# 按原索引+分组序号聚合，得到每组的(id,x,y)单行数据，过滤空组
df_long = df_long.pivot(index=['idx', 'group'], columns='type', values='value').reset_index()
df_long = df_long.dropna(subset=['id']).astype({'id': int})

# 2. 以id为列、x/y为值转回宽格式，构造目标列名
df_result = df_long.pivot(index='idx', columns='id', values=['x', 'y'])
df_result.columns = [f'{col[0]}_{col[1]}' for col in df_result.columns]

# 3. 调整列顺序与目标一致（可选）
target_cols = ['x_2', 'y_2', 'x_4', 'y_4', 'x_5', 'y_5', 'x_6', 'y_6', 'x_9', 'y_9']
df_result = df_result[target_cols].reset_index(drop=True)

# 验证结果是否匹配
print(df_result.equals(df_after))  # 输出 True

关键步骤说明 #

• unstack()将原宽表的列拆分为行，保留原索引（idx）确保后续能对应回原始行数据；
• 拆分列名后聚合，过滤掉无有效id的空组，得到干净的(idx, id, x, y)长表；
• 二次pivot将id转为列，合并层级列名得到x_id/y_id的格式，最后调整列顺序即可匹配目标结构。

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