如何程序化实现AB测试用户组多维度均匀拆分及销售人均客户数均衡？

核心需求 #

• 按固定比例将每个销售负责人（sales_person）的客户拆分到Treatment组和Control组
• 最小化两组在age、prior_sales、income三个指标上的均值差异

方案一：整数二次规划（最优解） #

使用线性规划库pulp构建优化模型，在满足销售负责人比例约束的前提下，最小化两组均值的平方差之和，得到理论最优分组。

步骤说明 #

1. 计算每个销售负责人需分配到Treatment组的客户数量（确保为整数）
2. 定义优化目标：最小化Treatment组与整体均值的平方差之和
3. 添加约束：每个销售负责人的Treatment组客户数严格符合指定比例
4. 求解模型并完成分组分配

代码实现 #

import pandas as pd
import pulp

# 加载示例数据
data = pd.DataFrame({
    'customer_id': [1,2,3,4,5,6,7,8],
    'age': [34,42,22,34,25,64,27,49],
    'prior_sales': [500,600,300,400,300,150,200,750],
    'income': [50000,90000,70000,95000,45000,125000,85000,105000],
    'sales_person': ['Tom','Sarah','Tom','Tom','Sarah','Jerry','Sarah','Jerry']
})

# 参数设置
split_ratio = 0.5
balance_vars = ['age', 'prior_sales', 'income']

# 计算每个销售负责人的Treatment组客户数
sales_person_counts = data['sales_person'].value_counts()
n_treat_per_sp = (sales_person_counts * split_ratio).astype(int)
# 若比例导致非整数，可根据实际情况微调（如四舍五入或全局调整）

# 计算整体均值
overall_means = data[balance_vars].mean()
total_treat = n_treat_per_sp.sum()

# 初始化优化问题
prob = pulp.LpProblem("CustomerSplitOptimization", pulp.LpMinimize)

# 创建二进制变量：1=Treatment组，0=Control组
customer_vars = pulp.LpVariable.dicts("Customer", data.index, cat='Binary')

# 定义目标函数：最小化Treatment组与整体均值的平方差之和
objective = 0
for var in balance_vars:
    sum_treat = pulp.lpSum([customer_vars[i] * data.loc[i, var] for i in data.index])
    mean_diff = (sum_treat / total_treat) - overall_means[var]
    objective += mean_diff ** 2

prob += objective

# 添加约束：每个销售负责人的Treatment组客户数符合比例
for sp in sales_person_counts.index:
    sp_customers = data[data['sales_person'] == sp].index
    prob += pulp.lpSum([customer_vars[i] for i in sp_customers]) == n_treat_per_sp[sp]

# 求解模型（需支持二次规划的Solver，如Gurobi/CPLEX，或开源SCIP）
prob.solve(pulp.GUROBI(msg=0))

# 分配分组
data['Group'] = data.index.map(lambda i: 'Treatment' if pulp.value(customer_vars[i]) == 1 else 'Control')

# 输出结果对比
print("整体均值:\n", overall_means)
print("\nTreatment组均值:\n", data[data['Group'] == 'Treatment'][balance_vars].mean())
print("\nControl组均值:\n", data[data['Group'] == 'Control'][balance_vars].mean())

方案二：遗传算法（启发式解） #

当数据量较大时，使用遗传算法寻找近似最优解，避免整数规划的计算复杂度，适合大规模数据集。

步骤说明 #

1. 生成初始种群：每个个体是符合销售负责人比例约束的分组方案
2. 定义适应度函数：以均值平方差之和作为优化目标
3. 变异操作：在同一销售负责人的客户中交换Treatment/Control组的客户
4. 迭代进化，筛选最优分组方案

代码实现 #

import pandas as pd
import numpy as np
from deap import base, creator, tools, algorithms
import random

# 加载示例数据
data = pd.DataFrame({
    'customer_id': [1,2,3,4,5,6,7,8],
    'age': [34,42,22,34,25,64,27,49],
    'prior_sales': [500,600,300,400,300,150,200,750],
    'income': [50000,90000,70000,95000,45000,125000,85000,105000],
    'sales_person': ['Tom','Sarah','Tom','Tom','Sarah','Jerry','Sarah','Jerry']
})

# 参数设置
split_ratio = 0.5
balance_vars = ['age', 'prior_sales', 'income']
pop_size = 50
num_generations = 100
mutation_prob = 0.5

# 计算每个销售负责人的Treatment组客户数
sales_person_counts = data['sales_person'].value_counts()
n_treat_per_sp = (sales_person_counts * split_ratio).astype(int)
overall_means = data[balance_vars].mean()

# 定义适应度函数：最小化均值平方差之和
def evaluate(individual):
    treat_mask = np.array(individual) == 1
    treat_means = data.loc[treat_mask, balance_vars].mean()
    diff_sum = ((treat_means - overall_means) ** 2).sum()
    return (diff_sum,)

# 初始化遗传算法组件
creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
creator.create("Individual", list, fitness=creator.FitnessMin)

toolbox = base.Toolbox()

# 生成符合约束的初始个体
def create_individual():
    individual = [0] * len(data)
    for sp in sales_person_counts.index:
        sp_indices = data[data['sales_person'] == sp].index.tolist()
        treat_indices = random.sample(sp_indices, n_treat_per_sp[sp])
        for idx in treat_indices:
            individual[idx] = 1
    return creator.Individual(individual)

toolbox.register("individual", create_individual)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)

# 变异操作：交换同一销售负责人下的Treatment/Control客户
def mutate(individual):
    sp = random.choice(sales_person_counts.index)
    sp_indices = data[data['sales_person'] == sp].index.tolist()
    treat_indices = [i for i in sp_indices if individual[i] == 1]
    control_indices = [i for i in sp_indices if individual[i] == 0]
    if not treat_indices or not control_indices:
        return (individual,)
    t_idx = random.choice(treat_indices)
    c_idx = random.choice(control_indices)
    individual[t_idx], individual[c_idx] = 0, 1
    return (individual,)

toolbox.register("mutate", mutate)
toolbox.register("select", tools.selTournament, tournsize=3)
toolbox.register("evaluate", evaluate)

# 运行遗传算法
pop = toolbox.population(n=pop_size)
hof = tools.HallOfFame(1)
stats = tools.Statistics(lambda ind: ind.fitness.values)
stats.register("avg", np.mean)
stats.register("min", np.min)

pop, log = algorithms.eaSimple(pop, toolbox, cxpb=0.0, mutpb=mutation_prob, 
                               ngen=num_generations, stats=stats, halloffame=hof, verbose=True)

# 分配最优分组
best_ind = hof[0]
data['Group'] = ['Treatment' if x == 1 else 'Control' for x in best_ind]

# 输出结果对比
print("整体均值:\n", overall_means)
print("\nTreatment组均值:\n", data[data['Group'] == 'Treatment'][balance_vars].mean())
print("\nControl组均值:\n", data[data['Group'] == 'Control'][balance_vars].mean())

方案选择 #

• 整数规划：适合数据量较小的场景，能得到理论最优解，但需要支持二次规划的求解器
• 遗传算法：适合大规模数据集，通过迭代寻找近似最优解，无需特殊求解器，计算效率更高

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