如何高效基于DataFrame概率列生成选择列？优化随机调用

首先纠正你原代码中的错误：random.choices的权重参数里，选选项2的概率应该是1-p而非p-1，否则当p<1时会出现负权重，导致报错。正确的单行写法应为：

df['choice'] = df['p'].apply(lambda p: random.choices(population=[1, 2], weights=[p, 1-p], k=1)[0])

但这种逐行调用random.choices的方式在大数据量或多次模拟时效率极低，结合你提到的p为离散值且大量重复的特点，推荐以下两种高效方案：

方案1：Numpy向量化生成（最通用高效） #

利用Numpy的向量化随机函数直接批量生成结果，避免Python循环，速度提升非常明显：

import numpy as np
import pandas as pd

# 示例数据
df = pd.DataFrame({'id': ['A', 'C', 'B', 'E', 'D'], 'p': [0.2, 0.4, 0.7, 0.2, 0.9]})

# 单次模拟生成choice列
df['choice'] = np.where(np.random.binomial(n=1, p=df['p'], size=len(df)) == 1, 1, 2)

# 多次模拟生成多列（比如模拟5次）
num_simulations = 5
for sim_idx in range(num_simulations):
    df[f'choice_{sim_idx+1}'] = np.where(np.random.binomial(n=1, p=df['p'], size=len(df)) == 1, 1, 2)

原理：np.random.binomial(n=1, p=p_val)本质是伯努利试验，生成1的概率为p_val（对应选选项1），生成0的概率为1-p_val（对应选选项2），最后用np.where把0替换为2即可。

方案2：基于离散p值的分组优化（极端重复场景下进一步提速） #

如果p的唯一值极少且每个值对应的样本量极大，可以先为每个唯一p批量生成对应数量的随机结果，再映射回原DataFrame，减少重复的随机数生成调用：

import numpy as np
import pandas as pd

df = pd.DataFrame({'id': ['A', 'C', 'B', 'E', 'D'], 'p': [0.2, 0.4, 0.7, 0.2, 0.9]})

# 获取每个p值对应的样本数量
p_counts = df['p'].value_counts()
# 为每个唯一p预先生成足够的随机结果
p_results = {}
for p_val in p_counts.index:
    count = p_counts[p_val]
    # 生成count个选择结果
    results = np.where(np.random.binomial(n=1, p=p_val, size=count) == 1, 1, 2)
    p_results[p_val] = results

# 通过groupby将结果分配回原数据
df['choice'] = df.groupby('p')['p'].transform(lambda x: p_results[x.name][:len(x)])

两种方案对比：方案1实现简单，在绝大多数场景下速度已经足够快；方案2在p的唯一值极少（比如仅2-3种）且单种p对应十万级以上样本时，能进一步减少随机函数的调用开销，但实现稍复杂。

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