蓄水池抽样算法

发表于 | 分类于 | | 阅读次数: |
字数统计: 660 | 阅读时长 ≈ 2

问题描述

现在有一组数，不知道这组数的总量有多少，请描述一种算法能够在这组数据中随机抽取k个数，使得每个数被取出来的概率相等。

蓄水池算法原理

随机抽样需要实现在n个元素中取k个元素时，每个元素被取中的概率都是相同的，即蓄水池算法可以在不知道n为多大的情况下保证每个元素被取中的概率相同。假设i是蓄水池算法取样的第i个元素，推导如下：

  • 对于第i个数（i<=k)，元素直接进入池中，每个元素进入池子的概率都是1

    • 第k+1个元素不被放进池子的概率=1-第k+1个元素被选中的概率$\times$第i个元素被选中的概率 = $1 - \frac{k}{k+1}*\frac{1}{k} = 1- \frac{1}{k+1} = \frac{k}{k+1}$
    • 第k+2个元素不被选中放进池子的概率=1-第k+2个元素被选中的概率$\times$第i个元素被选中的概率 = $1 - \frac{k}{k+2}*\frac{1}{k} = 1 - \frac{1}{k+2} = \frac{k+1}{k+2}$
    • …
    • 则运行到第n步时，第i个数仍保留的概率 = 被选中的概率 * 不被替换的概率，即：

  • 对于第j个数（j>k），在第j步被选中的概率为$\frac{k}{j}$

    • 不被第j+1个元素替换的概率为：$1 - \frac{k}{j + 1} \times \frac{1}{k} = \frac{j}{j + 1}$
    • 不被第j+2个元素替换的概率为：$1 - \frac{k}{j + 2} \times \frac{1}{k} = \frac{j+1}{j + 2}$
    • …
    • 则运行到第 n 步时，被保留的概率 = 被选中的概率 * 不被替换的概率，即：

所以对于其中每个元素，被保留的概率都为$\frac{k}{n}$

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import random

class ReservoirSampling:

    def __init__(self,  pools):
        self.pools = pools

    def sampling(self, k):
        assert k <= len(self.pools), 'k should be less than n'
        res = []
        # 前 K 个元素直接放入蓄水池中
        for i in range(k):
            res.append(self.pools[i])
        
        # K+1个元素开始进行概率采样，如果采样的r小于k，则用当前的j元素替换蓄水池中的r元素
        for j in range(k, len(self.pools)):
            r = random.randint(0, j)
            if r<k:
                res[r] = self.pools[j]
        return res

# 测试
res = []
for i in range(100000):
    k = 5
    N = 100
    pools = list(range(1, N+1))
    rs = ReservoirSampling(pools)
    s_res = rs.sampling(k)
    res += s_res

from collections import Counter

d = Counter(res)
print(d)