如何将数据集按value列均值指定比例划分为A/B/C三类标签

核心逻辑 #

我们以B类均值为基准值，根据输入的目标比例（如A:B:C=1.5:1:1.8）将value列的记录分配到三个分类中，通过「预分配-校验-微调」的流程保证最终均值比例误差控制在可接受范围内。

具体实现步骤（Python Pandas环境） #

前置准备 #

首先导入依赖包，构造测试数据，定义目标比例：

import pandas as pd
import numpy as np

# 构造示例数据集，替换为你自己的数据集即可
df = pd.DataFrame({'value': np.random.randint(10, 100, size=1000)})
# 定义目标比例，顺序为A、B、C的均值相对于B的比例，可按需修改
target_ratio = [1.5, 1, 1.8]

步骤1：预分配分类 #

先对value排序，按照比例对应的均值高低初步分配分类，比例越高的分类初始分配的value区间越高，能大幅提升后续迭代收敛速度：

# 先对value升序排序
df_sorted = df.sort_values('value', ignore_index=True)
n = len(df_sorted)
# 初始按三等分预分配，可根据实际目标比例调整分桶大小
split1, split2 = int(n/3), int(2*n/3)
# 示例比例中B均值最低、其次是A、C最高，所以按value从低到高分给B、A、C
df_sorted.loc[:split1, 'category'] = 'B'
df_sorted.loc[split1:split2, 'category'] = 'A'
df_sorted.loc[split2:, 'category'] = 'C'

步骤2：迭代微调使均值比例符合要求 #

通过少量调整边界上的记录，把实际比例和目标比例的误差降到预设范围内：

# 允许的最大比例误差，要求越高数值设得越小
tolerance = 0.01
# 最大迭代次数，避免极端情况死循环
max_iter = 100

for _ in range(max_iter):
    # 计算当前三类的均值
    mean_a = df_sorted[df_sorted['category'] == 'A']['value'].mean()
    mean_b = df_sorted[df_sorted['category'] == 'B']['value'].mean()
    mean_c = df_sorted[df_sorted['category'] == 'C']['value'].mean()
    # 计算当前相对于B的比例
    current_ratio_a = mean_a / mean_b
    current_ratio_c = mean_c / mean_b
    # 满足误差要求则停止迭代
    if abs(current_ratio_a - target_ratio[0]) < tolerance and abs(current_ratio_c - target_ratio[2]) < tolerance:
        break
    # 调整A和B的边界
    if current_ratio_a < target_ratio[0]:
        # A的均值不够高，把B区最高的记录移到A区
        b_max_idx = df_sorted[df_sorted['category'] == 'B']['value'].idxmax()
        df_sorted.loc[b_max_idx, 'category'] = 'A'
    elif current_ratio_a > target_ratio[0]:
        # A的均值太高，把A区最低的记录移到B区
        a_min_idx = df_sorted[df_sorted['category'] == 'A']['value'].idxmin()
        df_sorted.loc[a_min_idx, 'category'] = 'B'
    # 调整A和C的边界
    if current_ratio_c < target_ratio[2]:
        # C的均值不够高，把A区最高的记录移到C区
        a_max_idx = df_sorted[df_sorted['category'] == 'A']['value'].idxmax()
        df_sorted.loc[a_max_idx, 'category'] = 'C'
    elif current_ratio_c > target_ratio[2]:
        # C的均值太高，把C区最低的记录移到A区
        c_min_idx = df_sorted[df_sorted['category'] == 'C']['value'].idxmin()
        df_sorted.loc[c_min_idx, 'category'] = 'A'

# 如果需要保留原始数据集的顺序，可提前记录原始index，此处为打乱顺序输出示例
df = df_sorted.sample(frac=1, ignore_index=True)

步骤3：结果校验 #

# 输出最终三类的均值比例
final_means = df.groupby('category')['value'].mean()
print("A/B均值比：", final_means['A']/final_means['B'])
print("C/B均值比：", final_means['C']/final_means['B'])
# 输出前几行结果
print(df.head())

注意事项 #

• 如果数据集样本量过小（小于100条），可能无法达到过高的精度要求，可以适当调大tolerance参数
• 目标比例差异过大时（比如单类均值是其他类的3倍以上），可以增加max_iter参数或者调整初始分桶的比例
• 如果目标比例的高低顺序发生变化，对应调整初始预分配的分类和value区间的对应关系即可

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