使用Pandas构建复杂透视表并添加rem级别行求和列及代码优化

一、简化多步透视表流程 #

原四次独立pivot_table加concat的写法冗余，可通过单次透视+列层级处理直接生成目标结构：

假设原DataFrame为df，type取值为1-4，代码如下：

# 单次生成含type维度的透视表
pivot_df = pd.pivot_table(
    df,
    index=['rem', 'rp', 'road'],  # 保留原透视的行索引
    columns='type',
    values='overall',
    aggfunc='sum'  # 按实际需求替换聚合函数，如'mean'
)

# 重命名列得到overall_1至overall_4
pivot_df.columns = [f'overall_{t}' for t in pivot_df.columns]
pivot_df = pivot_df.reset_index()  # 可选：将索引转为普通列，按需调整

此方法避免了重复调用API，代码更简洁且执行效率更高。

二、添加rem=2/3的分组求和列 #

根据需求，可选择将求和结果作为新列或新行添加：

方式1：求和结果作为新列 #

# 筛选rem=2/3的行，按rp、road分组求和
sum_2 = pivot_df[pivot_df['rem'] == 2].groupby(['rp', 'road'])[[f'overall_{t}' for t in range(1,5)]].sum()
sum_3 = pivot_df[pivot_df['rem'] == 3].groupby(['rp', 'road'])[[f'overall_{t}' for t in range(1,5)]].sum()

# 重命名求和列
sum_2.columns = [col.replace('overall', '2_sum') for col in sum_2.columns]
sum_3.columns = [col.replace('overall', '3_sum') for col in sum_3.columns]

# 合并回原透视表
final_df = pivot_df.merge(sum_2, on=['rp', 'road'], how='left')
final_df = final_df.merge(sum_3, on=['rp', 'road'], how='left')

方式2：求和结果作为新行 #

# 生成rem=2的求和行，标记rem为'2_sum'
sum_row_2 = pivot_df[pivot_df['rem'] == 2].groupby(['rp', 'road'])[[f'overall_{t}' for t in range(1,5)]].sum()
sum_row_2['rem'] = '2_sum'

# 生成rem=3的求和行，标记rem为'3_sum'
sum_row_3 = pivot_df[pivot_df['rem'] == 3].groupby(['rp', 'road'])[[f'overall_{t}' for t in range(1,5)]].sum()
sum_row_3['rem'] = '3_sum'

# 合并原表与求和行
final_df = pd.concat([pivot_df, sum_row_2.reset_index(), sum_row_3.reset_index()], ignore_index=True)

完整示例代码 #

import pandas as pd

# 模拟测试数据
data = {
    'rem': [1,2,3,1,2,3,1,2,3],
    'rp': ['A','A','A','B','B','B','C','C','C'],
    'road': ['X','X','X','X','X','X','Y','Y','Y'],
    'type': [1,1,1,2,2,2,3,3,3],
    'overall': [10,20,30,15,25,35,5,15,25]
}
df = pd.DataFrame(data)

# 生成优化后的透视表
pivot_df = pd.pivot_table(
    df,
    index=['rem', 'rp', 'road'],
    columns='type',
    values='overall',
    aggfunc='sum'
)
pivot_df.columns = [f'overall_{t}' for t in pivot_df.columns]
pivot_df = pivot_df.reset_index()

# 添加列形式求和列
sum_2 = pivot_df[pivot_df['rem'] == 2].groupby(['rp', 'road'])[[f'overall_{t}' for t in range(1,4)]].sum()
sum_2.columns = [col.replace('overall', '2_sum') for col in sum_2.columns]
sum_3 = pivot_df[pivot_df['rem'] == 3].groupby(['rp', 'road'])[[f'overall_{t}' for t in range(1,4)]].sum()
sum_3.columns = [col.replace('overall', '3_sum') for col in sum_3.columns]

final_df = pivot_df.merge(sum_2, on=['rp', 'road'], how='left').merge(sum_3, on=['rp', 'road'], how='left')
print(final_df)

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