如何按概率分布从列表中选取多个不重复的元素？

问题背景 #

random.sample(population, k) 能满足无放回抽样的需求，但不支持传入概率权重；random.choices(population, weights, k) 支持权重设置，却可能返回重复元素，比如：

>>> random.choices([1, 2], [0.5, 0.5], k=2)
[1, 1] 

而我们期望的输出始终是 [1, 2] 或 [2, 1]。

解决方案1：纯Python自定义带权重无放回抽样 #

通过每次选择后移除已选元素及对应权重的方式，实现符合权重分布的无放回抽样：

import random

def weighted_sample(population, weights, k):
    # 复制原列表，避免修改输入数据
    remaining_pop = population.copy()
    remaining_weights = weights.copy()
    result = []
    
    for _ in range(k):
        # 基于当前剩余权重选择一个元素
        selected = random.choices(remaining_pop, weights=remaining_weights, k=1)[0]
        # 移除已选元素和对应权重
        idx = remaining_pop.index(selected)
        result.append(selected)
        remaining_pop.pop(idx)
        remaining_weights.pop(idx)
    
    return result

测试效果：

>>> weighted_sample([1, 2], [0.5, 0.5], k=2)
[2, 1]  # 或 [1, 2]，绝不会出现重复元素

解决方案2：用numpy实现高效抽样 #

如果处理的数据集规模较大，推荐使用numpy的原生方法，它更高效且代码更简洁：

import numpy as np

def weighted_sample_np(population, weights, k):
    # replace=False 表示无放回抽样，p 参数传入权重数组
    return np.random.choice(population, size=k, replace=False, p=weights).tolist()

测试效果：

>>> weighted_sample_np([1, 2], [0.5, 0.5], k=2)
[1, 2]  # 或 [2, 1]

重要提示 #

• 不要用random.choices抽样后再去重：这种操作会破坏原本的概率分布，而且当k等于列表长度时，去重后的结果可能无法满足数量要求。
• 权重数组不需要提前归一化（总和不必为1），两个方法都会自动处理权重的归一化。

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