问题描述
现在有一组数,不知道这组数的总量有多少,请描述一种算法能够在这组数据中随机抽取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 | import random |