如何重塑DataFrame计算指定日期与分组下两组采购量的差值？

问题描述 #

现有按date-bucket-group维度拆分的DataFrame：

date       | bucket |  Group  |purchase
2020-01-01 | 1      |  A      | 12
2020-01-01 | 1      |  B      | 11

2020-01-01 | 2      |  A      | 14
2020-01-01 | 2      |  B      | 14

2020-02-01 | 1      |  A      | 11
2020-02-01 | 1      |  B      | 10

需要生成包含difference列的新DataFrame，格式如下：

date       | bucket |  purchase | difference
2020-01-01 | 1      |  12-11=1  | Group A - Group B for that day/bucket
2020-01-01 | 2      |  0        | Group A - Group B for that day/bucket
2020-02-01 | 1      |  1        | Group A - Group B for that day/bucket

实现步骤 #

使用Pandas可以通过以下步骤快速实现：

1. 数据透视转换 #

将Group列的A、B分组转为独立列，方便后续计算：

import pandas as pd

# 假设原始数据存储在df中
pivoted = df.pivot(index=['date', 'bucket'], columns='Group', values='purchase').reset_index()

如果存在date-bucket组合下缺失A/B组的情况，添加fill_value=0避免空值：

pivoted = df.pivot(index=['date', 'bucket'], columns='Group', values='purchase').fillna(0).reset_index()

2. 计算差值与格式化显示列 #

# 计算A组减B组的差值
pivoted['difference'] = pivoted['A'] - pivoted['B']

# 格式化purchase列：差值非0时显示"X-Y=Z"，否则显示0
pivoted['purchase'] = pivoted.apply(
    lambda x: f"{x['A']}-{x['B']}={x['difference']}" if x['difference'] != 0 else '0',
    axis=1
)

3. 整理最终结构 #

保留所需列并调整顺序，同时统一difference列的说明文本：

final_df = pivoted[['date', 'bucket', 'purchase', 'difference']]
final_df['difference'] = 'Group A - Group B for that day/bucket'

运行上述代码后，final_df即为符合需求的结果。

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