如何在Pandas中按分段计算变体组相对对照组的%lift？

问题背景 #

原始DataFrame结构如下：

segment    group        purchase_amount 
A          control      2601    
A          variant1     2608    
A          variant2     2586    
B          control      2441    
B          variant1     2712    
B          variant2     2710    

需要生成表格，在各segment内将所有变体组与对应对照组对比，得到如下结果：

segment    %lift        purchase_amount 
A          variant1     (2608-2601)/2601
A          variant2     (2586-2601)/2601    
B          variant1     (2712-2441)/2441    
B          variant2     (2710-2441)/2441

实现方案 #

步骤说明 #

1. 提取对照组基准值：按segment分组，提取每个分段中group为control的purchase_amount，作为该分段的对比基准。
2. 合并基准值到原数据：将基准值合并回原始DataFrame，让每个变体组都能匹配到对应分段的对照组数值。
3. 过滤并计算提升率：移除对照组本身的行，按需生成提升率的公式字符串或实际数值，最后整理列名和结构。

代码实现 #

import pandas as pd

# 构造原始数据
df = pd.DataFrame({
    'segment': ['A', 'A', 'A', 'B', 'B', 'B'],
    'group': ['control', 'variant1', 'variant2', 'control', 'variant1', 'variant2'],
    'purchase_amount': [2601, 2608, 2586, 2441, 2712, 2710]
})

# 获取每个分段的对照组金额
control_data = df[df['group'] == 'control'].set_index('segment')['purchase_amount'].rename('control_amount')

# 合并对照组金额到原数据表
merged_df = df.merge(control_data, on='segment')

# 过滤掉对照组行，处理提升率
result_df = merged_df[merged_df['group'] != 'control'].copy()

# 方式1：生成公式形式的字符串（符合需求格式）
result_df['purchase_amount'] = result_df.apply(
    lambda row: f"({row['purchase_amount']}-{row['control_amount']})/{row['control_amount']}",
    axis=1
)

# 方式2：计算实际提升率数值（如果需要具体百分比，可打开下面注释）
# result_df['purchase_amount'] = (result_df['purchase_amount'] - result_df['control_amount']) / result_df['control_amount']

# 调整列名并删除冗余列
result_df = result_df.rename(columns={'group': '%lift'}).drop('control_amount', axis=1)

# 输出结果
print(result_df)

运行结果 #

执行代码后会输出符合要求的表格：

segment     %lift          purchase_amount
1       A  variant1        (2608-2601)/2601
2       A  variant2        (2586-2601)/2601
4       B  variant1        (2712-2441)/2441
5       B  variant2        (2710-2441)/2441

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