Pandas中如何高效按列分组将多行数据合并为目标宽表

需求说明 #

源数据按A列分组后，需要保留每组第1条记录的B、C列原值，将每组第2条记录的B、C值映射到新增D、E列，无对应记录的位置填充NaN，此前使用pd.merge、rank未找到高效实现路径。

源数据示例 #

目标结果示例 #

实现方案 #

用分组计数+透视的向量化方案实现，性能远高于merge类方案，操作步骤如下：

• 第一步：为每个A分组内的行生成从0开始的递增序号，标记行在组内的位置
• 第二步：过滤掉每组序号大于等于2的冗余记录，只保留需要转置的前2行
• 第三步：以A列为索引、组内序号为列做透视，将B、C列的值按组展开
• 第四步：重命名透视后的多级列名，调整列顺序后重置索引得到最终结果

对应可直接运行的代码：

import pandas as pd

# 构造示例源数据
df1 = pd.DataFrame({
    'A': ['TBK1', 'TBK1', 'TBK3'],
    'B': ['2022-01-01', '2022-02-02', '2022-05-07'],
    'C': ['2022-04-04', '2021-01-09', '2023-02-04']
})

# 生成组内行号
df1['group_row_idx'] = df1.groupby('A').cumcount()
# 过滤保留每组前2条
df1 = df1[df1['group_row_idx'] < 2]
# 透视展开列
pivot_res = df1.pivot(index='A', columns='group_row_idx', values=['B', 'C'])
# 列重命名+顺序调整
pivot_res.columns = ['B', 'D', 'C', 'E']
df2 = pivot_res[['B', 'C', 'D', 'E']].reset_index()

该方案全程使用Pandas内置向量化操作，无逐行循环、无多表关联的额外开销，在十万级以上数据量下运行效率比merge方案高一个数量级。如果后续需要保留每组更多行的字段，只需要调整列名映射规则和过滤阈值即可适配。

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