基于另一数据集概率抽取数据并分配Group字段的技术实现请求

数据集说明 #

数据集1：组别概率表 #

记录不同周期（Period）、社会人口统计组合（socio_demo）下，每个组别（Group）的分配概率（Probabilities，百分比），Group固定为1-4，共11个周期、144种socio_demo组合。示例数据：

数据集2：样本数据表 #

包含周期（Period）、样本ID、社会人口统计组合（socio_demo）、权重（Weight）字段，示例数据：

需求说明 #

针对数据集1的每个（Period, socio_demo）组合，按照各组别的概率百分比，从数据集2的对应（Period, socio_demo）样本中随机分配Group标签，确保每个组合下的样本按概率比例分到1-4组，最终输出包含数据集2所有字段+新增Group字段的完整表。

示例输出（对应上述示例数据）：

技术实现方案（Python Pandas） #

步骤1：数据导入与预处理 #

import pandas as pd
import numpy as np

# 导入两个数据集
df_prob = pd.read_csv("probability_dataset.csv")  # 替换为你的数据集路径
df_samples = pd.read_csv("sample_dataset.csv")    # 替换为你的数据集路径

# 校验每个（Period, socio_demo）组合的概率总和是否为100
grouped_prob = df_prob.groupby(["Period", "socio_demo"])["Probabilities"].sum()
assert (grouped_prob == 100).all(), "部分（Period, socio_demo）组合的概率和不为100，请检查数据"

步骤2：按分组分配Group标签 #

def assign_group(group):
    # 获取当前分组的周期和社会人口组合
    period = group["Period"].iloc[0]
    socio = group["socio_demo"].iloc[0]
    
    # 提取对应概率数据
    prob_data = df_prob[(df_prob["Period"] == period) & (df_prob["socio_demo"] == socio)]
    groups = prob_data["Group"].tolist()
    probs = [p / 100 for p in prob_data["Probabilities"].tolist()]  # 转为0-1比例
    
    # 随机分配Group
    group["Group"] = np.random.choice(groups, size=len(group), p=probs)
    return group

# 按（Period, socio_demo）分组处理样本
df_result = df_samples.groupby(["Period", "socio_demo"], group_keys=False).apply(assign_group)

# 重置索引（可选）
df_result = df_result.reset_index(drop=True)

步骤3：输出结果 #

# 保存为CSV文件
df_result.to_csv("assigned_group_result.csv", index=False)

# 打印前几行查看结果
print(df_result.head())

关键提示 #

• 概率校验：必须确保每个（Period, socio_demo）组合下的概率总和为100，否则np.random.choice会报错。若存在数据误差，可将概率除以总和做归一化处理。
• 可复现性：若需要固定分配结果，可在代码开头添加np.random.seed(42)（数字可自定义），保证每次运行结果一致。
• 性能优化：针对超大规模样本，可考虑按比例抽取各组样本后拼接，但上述分组apply方式更简洁，能保证所有样本都被分配。
• 加权分配：如果需要按Weight字段加权分配Group，可修改assign_group函数，将权重与概率结合后传入np.random.choice的p参数。

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