python - 使用 group by 透视 pandas 数据框

Question

我有一个像这样的数据框:

id  sub_id  count
0   94  1
1   94  9
1   315 7
2   94  4
2   265 1


data = {'id': [0,1,1,2,2], 
     'sub_id': [94,94,315,94,265], 
     'count': [1,9,7,4,1]
    }
df = pd.DataFrame(data)

And I want it in the following form:
id sub_id1 count_sub_id1 sub_id2 count_sub_id2
0  94      1             NaN     NaN
1  94      9             315     7
2  94      4             265     1

注意:这里，每个 id 最多可以有两行，每行都有不同的 sub_id 及其计数。

我尝试了这个 df.pivot(index='id',columns='sub_id',values='count') 但这导致第二列中的所有行都以不同的方式展开列，而我只需要两列，并具有自定义名称，即。仅每组 id 存在的两行

Answer 1

尝试使用:

df_out = (df.set_index(['id', df.groupby('id').cumcount()+1])
            .unstack().sort_index(level=1, axis=1))

df_out.columns = [f'{i}{j}' if i == "sub_id" else f'{i}_sub_id{j}' 
                          for i, j in df_out.columns]

print(df_out.reset_index())

输出:

   id  count_sub_id1  sub_id1  count_sub_id2  sub_id2
0   0            1.0     94.0            NaN      NaN
1   1            9.0     94.0            7.0    315.0
2   2            4.0     94.0            1.0    265.0