Pandas长表转宽表进阶：基于关联索引的转换需求

现有如下格式的Pandas DataFrame：

import pandas as pd
import numpy as np

ex_years = [2016, 2016, 2016, 2017, 2017, 2017, 2017, 2018, 2018, 2018]
ex_xs = [3, 1, 3, 2, 2, 5, 2, 1, 4, 1]
ind = [1, 2, 3, 1, 2, 3, 4, 1, 2, 3]
prev_ind = [np.nan, np.nan, np.nan, np.nan, 2, np.nan, 3, 1, 1, 4]
prev_year = [np.nan, np.nan, np.nan, np.nan, 2016, np.nan, 2016, 2016, 2017, 2017]
df = pd.DataFrame({'Year':ex_years, 'ind':ind, 'prev_ind':prev_ind, 'prev_year':prev_year, 'x':ex_xs})

原始DataFrame展示：

Year  ind  prev_ind  prev_year  x
0  2016    1       NaN        NaN  3
1  2016    2       NaN        NaN  1
2  2016    3       NaN        NaN  3
3  2017    1       NaN        NaN  2
4  2017    2       2.0     2016.0  2
5  2017    3       NaN        NaN  5
6  2017    4       3.0     2016.0  2
7  2018    1       1.0     2016.0  1
8  2018    2       1.0     2017.0  4
9  2018    3       4.0     2017.0  1

需要将其转换为规整的宽表格式，观测值通过prev_ind和prev_year关联（即当前观测值在prev_year对应的索引为prev_ind），预期输出如下：

2016  2017  2018
0   3.0   NaN   1.0
1   NaN   2.0   4.0
2   1.0   2.0   NaN
3   NaN   5.0   NaN
4   3.0   2.0   1.0

核心难点在于：个体在不同年份的索引不固定，部分个体仅出现一年，部分存在年份缺失。

• 核心逻辑：给属于同一个体的所有观测值分配唯一的group_id——无关联的观测值（prev_ind为NaN）作为个体起点，有前驱的观测值通过prev_year和prev_ind回溯到前驱，继承同一个group_id。
• 先给所有无关联行分配初始group_id，再迭代处理剩余行，直到所有行完成group_id分配。
• 最后用pivot方法将长表转宽表，以group_id为行索引，年份为列填充x值。

# 给所有行分配group_id
df['group_id'] = np.nan
# 给无关联的根节点分配初始ID
root_rows = df['prev_ind'].isna()
df.loc[root_rows, 'group_id'] = range(root_rows.sum())

# 迭代处理有前驱的行，直到所有group_id都被填充
while df['group_id'].isna().any():
    for idx, row in df[df['group_id'].isna()].iterrows():
        # 找到前驱行的group_id
        prev_row = df[(df['Year'] == row['prev_year']) & (df['ind'] == row['prev_ind'])]
        if not prev_row.empty and not np.isnan(prev_row['group_id'].values[0]):
            df.loc[idx, 'group_id'] = prev_row['group_id'].values[0]

# 转换为宽表
wide_df = df.pivot(index='group_id', columns='Year', values='x').reset_index(drop=True)
print(wide_df)

运行代码后即可得到预期的宽表格式。

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