如何在Pandas透视表中添加多列差值并还原初始格式？

原始DataFrame #

import pandas as pd

df = pd.DataFrame({"A": ["foo", "foo", "foo", "foo", "foo",
                         "bar", "bar", "bar", "bar",'foo' ],
                   "B": ["one", "one", "one", "two", "two",
                         "one", "one", "two", "two", 'two'],
                   "C": ["small", "large", "large", "small",
                         "small", "large", "small", "small",
                         "large", 'large'],
                   "D": [1, 2, 2, 3, 3, 4, 5, 6, 7,8],
               })

执行透视表操作 #

table = pd.pivot_table(df, values='D', index=['A'],
                    columns=['B','C'])

得到的透视表结果 #

B   one             two
C   large   small   large   small
A               
bar   4      5       7        6
foo   2      1       8        3

需要解决两个问题：

• 如何为"one"和"two"分组添加large - small的差值（命名为diff），理想情况下使用aggfunc实现？
• 如何将包含差值的透视表重新转换为初始数据格式？

一、添加差值列（使用aggfunc实现） #

方式1：自定义聚合函数生成结果 #

直接通过分组+自定义聚合函数，一次性生成均值和差值：

def agg_with_diff(group):
    # 计算large和small的均值
    large_mean = group[group['C'] == 'large']['D'].mean()
    small_mean = group[group['C'] == 'small']['D'].mean()
    # 返回包含均值和差值的Series
    return pd.Series({
        'large': large_mean,
        'small': small_mean,
        'diff': large_mean - small_mean
    })

# 按A、B分组聚合，再将B转为列
result = df.groupby(['A', 'B']).apply(agg_with_diff).unstack('B')

生成的结果自动保持多级列结构：

one                two              
       large small diff large small diff
A                                       
bar      4.0   5.0 -1.0   7.0   6.0  1.0
foo      2.0   1.0  1.0   8.0   3.0  5.0

方式2：基于已有透视表追加差值列 #

如果已经生成了初始透视表table，可以直接对多级列计算差值：

# 遍历B的每个分组，计算large与small的差值
for b in table.columns.get_level_values('B').unique():
    table[(b, 'diff')] = table[(b, 'large')] - table[(b, 'small')]

# 按B层级排序列，让diff紧跟对应分组
table = table.sort_index(axis=1)

最终结果和方式1完全一致，操作更简洁。

二、转换回初始数据格式 #

使用melt方法将透视表转回长格式，步骤如下：

# 将索引A转为普通列
long_df = table.reset_index()
# 多级列转长格式，保留A作为标识列
long_df = long_df.melt(id_vars='A', var_name=['B', 'C'], value_name='D')
# 清理空值并重置索引
long_df = long_df.dropna().reset_index(drop=True)

得到的long_df与原始df结构一致，新增了C='diff'的行：

A    B      C    D
0  bar  one  large  4.0
1  foo  one  large  2.0
2  bar  one  small  5.0
3  foo  one  small  1.0
4  bar  one   diff -1.0
5  foo  one   diff  1.0
6  bar  two  large  7.0
7  foo  two  large  8.0
8  bar  two  small  6.0
9  foo  two  small  3.0
10 bar  two   diff  1.0
11 foo  two   diff  5.0

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