通过创建新列扁平化Pandas DataFrame，实现唯一(id,sid)对

我有如下Pandas DataFrame：

id sid X_animal X_class Y_animal Y_class
0   1   A       88    Home   Monkey  Mammal
1   1   A       88    Home   Parrot    Bird
2   1   B
3   2   C       11    Work
4   2   C       11    Work
5   2   C       33  School      Dog  Mammal
6   3   D       44    Home   Salmon    Fish
7   3   D       44    Home     Bear  Mammal
8   3   D       44    Home      Dog  Mammal
9   4   E       55  School

希望将其扁平化，使得每行的(id,sid)对唯一。同一(id,sid)下，*_animal和*_class列值不同时，创建带序号的新列，目标DataFrame如下：

id sid X_animal_1 X_class_1 X_animal_2 X_class_2 Y_animal_1 Y_class_1 Y_animal_2 Y_class_2 Y_animal_3 Y_class_3
0   1   A         88      Home                          Monkey    Mammal     Parrot      Bird
1   1   B
2   2   C         11      Work         33    School        Dog    Mammal
3   3   D         44      Home                          Salmon      Fish       Bear    Mammal        Dog    Mammal
4   4   E         55    School

生成初始和目标DataFrame的代码：

import pandas as pd
from numpy import nan

cols = ['id', 'sid', 'X_animal', 'X_class', 'Y_animal', 'Y_class']
l = [
    [1, 'A', 88, 'Home', 'Monkey', 'Mammal'],
    [1, 'A', 88, 'Home', 'Parrot', 'Bird'],
    [1, 'B', nan, nan, nan, nan],
    [2, 'C', 11, 'Work', nan, nan],
    [2, 'C', 11, 'Work', nan, nan],
    [2, 'C', 33, 'School', 'Dog', 'Mammal'],
    [3, 'D', 44, 'Home', 'Salmon', 'Fish'],
    [3, 'D', 44, 'Home', 'Bear', 'Mammal'],
    [3, 'D', 44, 'Home', 'Dog', 'Mammal'],
    [4, 'E', 55, 'School', nan, nan],
]

df = pd.DataFrame(data=l, columns=cols)
print(df.fillna(''))

cols2 = ['id', 'sid', 'X_animal_1', 'X_class_1', 'X_animal_2', 'X_class_2', 'Y_animal_1', 'Y_class_1', 'Y_animal_2', 'Y_class_2', 'Y_animal_3', 'Y_class_3']
l2 = [
    [1, 'A', 88, 'Home', nan, nan, 'Monkey', 'Mammal', 'Parrot', 'Bird'],
    [1, 'B', nan, nan, nan, nan, nan, nan, nan, nan],
    [2, 'C', 11, 'Work', 33, 'School', 'Dog', 'Mammal', nan, nan],
    [3, 'D', 44, 'Home', nan, nan, 'Salmon', 'Fish', 'Bear', 'Mammal', 'Dog', 'Mammal'],
    [3, 'E', 55, 'School', nan, nan, nan, nan, nan, nan],
]

df2 = pd.DataFrame(data=l2, columns=cols2)
print(df2.fillna(''))

尝试过pivot()和pivot_table()但失败，因列数量不固定出现KeyError。

可以通过分组生成唯一序号、分别透视X/Y列组、合并结果的方式实现需求，代码如下：

import pandas as pd
import numpy as np

# 加载初始数据
cols = ['id', 'sid', 'X_animal', 'X_class', 'Y_animal', 'Y_class']
l = [
    [1, 'A', 88, 'Home', 'Monkey', 'Mammal'],
    [1, 'A', 88, 'Home', 'Parrot', 'Bird'],
    [1, 'B', np.nan, np.nan, np.nan, np.nan],
    [2, 'C', 11, 'Work', np.nan, np.nan],
    [2, 'C', 11, 'Work', np.nan, np.nan],
    [2, 'C', 33, 'School', 'Dog', 'Mammal'],
    [3, 'D', 44, 'Home', 'Salmon', 'Fish'],
    [3, 'D', 44, 'Home', 'Bear', 'Mammal'],
    [3, 'D', 44, 'Home', 'Dog', 'Mammal'],
    [4, 'E', 55, 'School', np.nan, np.nan],
]
df = pd.DataFrame(data=l, columns=cols)

# 1. 为每组内的X/Y列组合生成唯一序号
# X组：按id,sid分组，对X_animal+X_class的唯一组合分配序号
df['X_seq'] = df.groupby(['id', 'sid'])[['X_animal', 'X_class']].transform(
    lambda x: x.astype(str).agg('-'.join, axis=1)
).rank(method='dense').astype(int)

# Y组：按id,sid分组，对Y_animal+Y_class的唯一组合分配序号
df['Y_seq'] = df.groupby(['id', 'sid'])[['Y_animal', 'Y_class']].transform(
    lambda x: x.astype(str).agg('-'.join, axis=1)
).rank(method='dense').astype(int)

# 2. 分别对X、Y组进行透视，生成带序号的列
x_pivot = df.pivot(index=['id', 'sid'], columns='X_seq', values=['X_animal', 'X_class'])
x_pivot.columns = [f'{col[0]}_{col[1]}' for col in x_pivot.columns]
x_pivot = x_pivot.reset_index()

y_pivot = df.pivot(index=['id', 'sid'], columns='Y_seq', values=['Y_animal', 'Y_class'])
y_pivot.columns = [f'{col[0]}_{col[1]}' for col in y_pivot.columns]
y_pivot = y_pivot.reset_index()

# 3. 合并X和Y的透视结果，去重并填充空值
result = pd.merge(x_pivot, y_pivot, on=['id', 'sid'], how='left')
result = result.groupby(['id', 'sid']).first().reset_index()

# 输出与目标格式一致的结果
print(result.fillna(''))

关键步骤说明 #

• 生成序号：通过将每组内的X/Y列组合转为字符串，用rank(method='dense')生成连续序号，确保相同列组合对应同一个序号，不同组合分配新序号。
• 透视重组：分别对X、Y列组进行透视，将序号作为列名后缀，自动生成对应数量的带序号列，避免固定列数导致的KeyError。
• 合并去重：合并透视后的X/Y结果，通过分组取第一个非空值去除重复行，最终得到每行唯一的(id,sid)对。

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