Pandas按segment分组计算AB测试MannWhitney U检验p值的实现方法

完整实现代码如下：

import pandas as pd
import numpy as np
from scipy.stats import mannwhitneyu

# 示例数据构造
data = [['a', 'segment1', 12,14], ['a', 'segment1', 12,14], ['b', 'segment2', 12,11],['a', 'segment2', 10,11],
       ['b', 'segment1', 4,5], ['b', 'segment1', 32,15], ['b', 'segment2', 14,8],['a', 'segment2', 11,21],
       ['b', 'segment1', 1,21], ['b', 'segment1', 4,21], ['a', 'segment2', 6,32],['b', 'segment2', 3,21],
       ]
df_data = pd.DataFrame(data, columns = ['test_group', 'segment', 'feature1', 'feature2'])

# 存储各segment的检验p值
p_value_map = {}
# 按segment分组执行检验
for seg, seg_data in df_data.groupby('segment'):
    a_data = seg_data[seg_data['test_group'] == 'a']
    b_data = seg_data[seg_data['test_group'] == 'b']
    # 分别计算两个feature的检验p值
    f1_p = mannwhitneyu(a_data['feature1'], b_data['feature1'], use_continuity=True, alternative='two-sided').pvalue
    f2_p = mannwhitneyu(a_data['feature2'], b_data['feature2'], use_continuity=True, alternative='two-sided').pvalue
    p_value_map[seg] = {'feature1': f1_p, 'feature2': f2_p}

# 构造符合要求的输出表
result_df = df_data[['segment', 'test_group']].drop_duplicates().sort_values(['segment', 'test_group']).reset_index(drop=True)
# 填充p值，仅test_group为b的行保留结果
for i, row in result_df.iterrows():
    if row['test_group'] == 'b':
        result_df.loc[i, 'feature1'] = p_value_map[row['segment']]['feature1']
        result_df.loc[i, 'feature2'] = p_value_map[row['segment']]['feature2']
    else:
        result_df.loc[i, ['feature1', 'feature2']] = np.nan

# 若需要将空值替换为空白字符串，执行以下语句即可
# result_df = result_df.fillna('')

输出说明 #

运行代码得到的result_df完全符合你要求的结构，test_group为a的行对应feature字段默认填充NaN，可根据业务需要调用fillna('')替换为空字符串。

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