如何优化Pandas重复索引DataFrame的列转换代码

我有一个带重复索引的单列Pandas DataFrame：

value
1  0.996957
1  1.098198
1  1.184518
2  1.255916
2  1.312393

我想要得到一个以唯一索引为行索引的新DataFrame，把每个索引对应的不同值作为列，列数由出现次数最多的索引决定。期望结果如下：

1         2         3
1  0.996957  1.098198  1.184518
2  1.255916  1.312393       NaN

我目前的代码可以正常运行，但希望通过分组索引的方式优化它，现有代码如下：

data = pd.DataFrame(columns=np.arange(df.groupby(df.index).count().value.max()))
for i in np.unique(df.index.values):
    data_points_i = pd.DataFrame(df.loc[i,:].value.values).transpose()
    data = pd.concat([data,data_points_i], axis = 0)
data.index = df.index.unique()

其中df是原始DataFrame，data是最终生成的DataFrame，恳请帮助优化这段代码。

优化方案 #

你可以试试这种更简洁高效的方式，完全利用Pandas的分组和重塑函数，避免循环和多次concat操作：

# 给每个组内的元素生成从1开始的序号，作为后续的列名
df['col_idx'] = df.groupby(level=0).cumcount() + 1

# 将组内序号转为列，原索引保留为行索引
result = df.set_index('col_idx', append=True)['value'].unstack('col_idx')

运行后得到的result就是你想要的格式：

1         2         3
1  0.996957  1.098198  1.184518
2  1.255916  1.312393       NaN

代码逻辑说明 #

1. groupby(level=0).cumcount() + 1：按原索引（也就是level=0的索引）分组，给每个组内的元素从0开始计数，加1是为了让列名从1开始，和你的期望结果匹配。
2. set_index('col_idx', append=True)：把生成的组内序号作为二级索引，和原索引组合成多层索引。
3. unstack('col_idx')：将二级索引（组内序号）转成列，自动把每个原索引对应的多个值排布到同一行的不同列，元素数量不足的位置会自动填充NaN。

这种方法不仅代码更简洁，还能利用Pandas的内置优化提升性能，数据量越大，对比循环的优势越明显。

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