如何在Pandas中对DataFrame按Id前缀分组并横向扩展列？

原始数据代码 #

import pandas as pd

df = pd.DataFrame(
[
    ['7890-1', '12345N', 'John', 'Intermediate'],
    ['7890-4', '30909N', 'Greg', 'Intermediate'], 
    ['3300-1', '88117N', 'Mark', 'Advanced'],
    ['2502-2', '90288N', 'Olivia', 'Elementary'],
    ['7890-2', '22345N', 'Joe', 'Intermediate'],
    ['7890-3', '72245N', 'Ana', 'Elementary']
],
columns=['Id', 'Code', 'Person', 'Level'])

解决方案代码 #

# 拆分Id列，提取前缀作为新Id、序号用于排序
df[['Id_prefix', 'seq']] = df['Id'].str.split('-', expand=True)
df['seq'] = df['seq'].astype(int)

# 按新Id分组，根据序号排序
df_sorted = df.sort_values(['Id_prefix', 'seq'])

# 转成宽表，以新Id为索引、序号为列层级
wide_df = df_sorted.pivot(index='Id_prefix', columns='seq', values=['Code', 'Person', 'Level'])

# 合并列层级，改为"列名 序号"格式
wide_df.columns = [f'{col[0]} {col[1]}' for col in wide_df.columns]

# 重置索引并恢复Id列名
wide_df = wide_df.reset_index().rename(columns={'Id_prefix': 'Id'})

# 补充缺失列，确保每组都有4组Code/Person/Level列
max_seq = 4
for seq in range(1, max_seq+1):
    for col in ['Code', 'Person', 'Level']:
        col_name = f'{col} {seq}'
        if col_name not in wide_df.columns:
            wide_df[col_name] = pd.NA

# 按指定顺序排列列
column_order = ['Id'] + [f'{col} {seq}' for seq in range(1, max_seq+1) for col in ['Code', 'Person', 'Level']]
wide_df = wide_df[column_order]

print(wide_df)

代码说明 #

1. 拆分Id：通过str.split分割原始Id，分离出用于分组的前缀和用于排序的序号，同时把序号转为整数类型。
2. 排序分组数据：确保每个Id前缀组内的条目按序号从小到大排列，对应最终宽表的列顺序。
3. 转宽表：利用pivot将长表转为宽表，把每组内的多条数据横向展开。
4. 列名整理：将双层列名合并为需求中的Code 1、Person 1格式，同时补充缺失的列保证结构统一。
5. 列序调整：按照需求的列顺序重新排列，确保输出结构完全匹配预期。

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