如何在Pandas分层索引中计算各points层级下团队A与B的差值

问题背景 #

已有如下Pandas DataFrame及分组后的结果：

import pandas as pd

df = pd.DataFrame({'team' : ['A', 'B', 'A', 'B', 'A', 'B'],
                   'tiger' : [87, 159, 351, 140, 72, 119],
                   'lion' : [1843, 3721, 6905, 1667, 2865, 1599],
                   'bear' : [1.9, 3.3, 6.3, 2.3, 1.2, 4.1],
                   'points' : [425, 425, 441, 441, 1048, 1048]})

grouped = df.groupby(['points', 'team'])[['tiger', 'lion', 'bear']].median()

分组后输出：

tiger       lion    bear
points team                             
425    A     87.00000 1843.00000 1.90000
       B    159.00000 3721.00000 3.30000
441    A    351.00000 6905.00000 6.30000
       B    140.00000 1667.00000 2.30000
1048   A     72.00000 2865.00000 1.20000
       B    119.00000 1599.00000 4.10000

需要在每个points层级下，计算团队B与团队A在tiger、lion、bear列的差值，得到如下格式结果：

points  tiger  lion     bear
0     425     72  1878  1.40000
1     441   -211 -5238 -4.00000
2    1048     47 -1266  2.90000

解决方案 #

方法一：基于分层索引的unstack操作 #

通过unstack()将team索引转为列，直接对同指标的B、A列做差值计算：

# 将team索引转为列，得到宽表格式
unstacked = grouped.unstack('team')

# 计算B列与A列的差值
diff_result = unstacked.xs('B', level='team', axis=1) - unstacked.xs('A', level='team', axis=1)

# 重置索引并清理列名格式
diff_result = diff_result.reset_index().rename_axis(None, axis=1)

print(diff_result)

方法二：直接在分组时计算差值 #

跳过中间的grouped变量，直接按points分组后，每组内提取A、B团队的中位数做差：

def calculate_diff(group):
    # 获取组内A、B团队的中位数
    a_median = group[group['team'] == 'A'][['tiger', 'lion', 'bear']].median()
    b_median = group[group['team'] == 'B'][['tiger', 'lion', 'bear']].median()
    # 返回B-A的差值
    return b_median - a_median

# 按points分组并应用计算函数
diff_result = df.groupby('points').apply(calculate_diff).reset_index()

print(diff_result)

两种方法都能得到预期的输出结果。

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