跳至主要內容

Python实现蓄水池算法

bbruceyuan小于 1 分钟面试锦囊面试锦囊

蓄水池算法

蓄水池算法采样算法难的点不在于怎么实现蓄水池算法，难的点在于证明每一个点都能被同等的概率抽取。

蓄水池算法证明最好的两个教程是：

  1. https://www.cnblogs.com/snowInPluto/p/5996269.htmlopen in new window
  2. https://www.jianshu.com/p/7a9ea6ece2afopen in new window

其他的感觉说的不是很清楚。

蓄水池算法的主要逻辑： image.png

蓄水池算法的Pthon实现：

import random

def reservior_sampling(n, k):
    """
    表示有 n 个数，随机采样 k 个
    """
    nums = [i for i in range(1, n + 1)]

    res = []
    for i in range(k):
        # 前K个数字可以直接填充
        res.append(nums[i])
    
    for i in range(k, len(nums)):
        # 假设 i == k (也是说这是 第 k + 1个元素), 那么 该数字有  k / (k + 1) 的概率被选中去去换
        replace_idx = random.randint(0, i)
        if replace_idx < k:
            res[replace_idx] = nums[i]
    return res

pool = reservior_sampling(100, 10)
print(pool)
# [78, 52, 41, 84, 66, 43, 25, 71, 45, 24]