如何用Pandas .apply()按分类权重概率为另一DataFrame生成多列抽取ID

实现代码 #

import pandas as pd
import numpy as np

# 构造并预处理第一个DataFrame
df = pd.DataFrame({
    'id': list(np.arange(1, 12)),
    'category': ['a', 'a', 'a', 'a', 'b', 'b', 'c', 'c', 'c', 'c', 'c'],
    'weight': list(np.random.randint(1, 5, 11))
})
df['weight_sum'] = df.groupby('category')['weight'].transform('sum')
df['p'] = df['weight'] / df['weight_sum']

# 构造第二个DataFrame示例，可替换为你的实际数据
df2 = pd.DataFrame({
    'id': [2, 3, 4],
    'category': ['a', 'a', 'a']
})

# 预构建每个分类的候选池，避免重复计算提升效率
category_pool = df.groupby('category').agg(
    candidate_ids=('id', list),
    candidate_p=('p', list)
).to_dict('index')

# 定义抽取函数
def sample_3_ids(category):
    pool = category_pool[category]
    ids = pool['candidate_ids']
    p = np.array(pool['candidate_p'])
    # 处理浮点数精度问题，保证概率和严格为1
    p = p / p.sum()
    # 无重复抽取3个id
    return np.random.choice(ids, size=3, replace=False, p=p)

# 执行抽取并生成结果
df2[['id_1', 'id_2', 'id_3']] = df2['category'].apply(sample_3_ids).apply(pd.Series)
result = df2.rename(columns={'id': 'user_id'})[['user_id', 'id_1', 'id_2', 'id_3']]

# 输出结果
print(result)

注意事项 #

之前用np.random.choice运行失败，通常是两个原因导致的：

1. 浮点数精度问题导致p列求和不等于1，上面代码中额外做了一次归一化可以解决该问题
2. 部分分类下的候选id数量少于3个，无重复抽取会报错，需要提前校验第一个DataFrame的数据是否符合要求

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