19. 蓄水池抽样算法

面试谈的算法

1. 问题模型描述：

给出一个数据流，我们需要在此数据流中随机选取 k 个数。由于这个数据流的长度很大，因此需要边遍历边处理，而不能将其一次性全部加载到内存。
请写出一个随机选择算法，使得数据流中所有数据被等概率选中。

2. 算法流程描述：

1. 构建一个大小为N的数组，初始将数据流中的前N个元素放入数组中，N即蓄水池的大小。
2. 从数据流的第N+1个元素开始，在[1, i]之间随机生成一个数rand，其中i表示当前是第几个元素。
3. 规定之后的元素以 N/i 的概率被选中，即rand>N时表示当前数没有被选中，不做任何操作，反之，当前数被选中，随机替换蓄水池中的一个元素。
4. 最终返回蓄水池中的所有元素信息。

通过上述流程之后，能够保证对于数据流中的元素能够被等概率的抽取到。该算法的核心在于先以某一种概率选取数，并在后续过程以另一种概率换掉之前已经被选中的数。因此实际上每个数被最终选中的概率都是被选中的概率 * 不被替换的概率。

3. 举例说明：

假设蓄水池大小N=10，总共数据流元素为1729个，
（1）计算第1729号元素到来时，3号元素依旧存货的概率，即3号元素最终被选中的概率
10号元素之前，3号元素存活的概率为1，
11号元素到来时，3号元素被淘汰的概率计算公式为：11号元素被选中的概率3号元素被选中淘汰的概率
，进而可以得到11号元素到来时3号元素存货的概率P(11)为，
12号元素到来时，3号元素被淘汰的概率计算公式为：12号元素被选中的概率3号元素被选中淘汰的概率
，进而可以得到12号元素到来时3号元素存货的概率P(12)为，
13号元素到来时，3号元素被淘汰的概率计算公式为：13号元素被选中的概率3号元素被选中淘汰的概率
，进而可以得到12号元素到来时3号元素存货的概率P(13)为，
…….依次类推
1729号元素到来时，3号元素被淘汰的概率计算公式为：1729号元素被选中的概率3号元素被选中淘汰的概率，即，进而可以得到1729号元素到来时3号元素存活的概率P(1729)为
根据条件概率公式，当1729号元素到来时3号元素存活的概率为：
=第3步被选中概率 不被 11 号元素替换的概率 不被 12 替换的概率 … 不被 1729替换的概率，
等价于 1 (1 - 被 11 替换的概率) (1 - 被12 替换的概率) … (1 - 被 1729替换的概率),
等价于

（2）计算第1729号元素到来时，13号元素依旧存货的概率，即3号元素最终被选中的概率
13>10，当13号元素到来时，13号元素被选中的概率为
14号元素到来时，13号元素被淘汰的概率计算公式为：14号元素被选中的概率13号元素被选中淘汰的概率，进而可以得到14号元素到来时13号元素存活的概率P(14)为，
15号元素到来时，13号元素被淘汰的概率计算公式为：15号元素被选中的概率13号元素被选中淘汰的概率，进而可以得到15号元素到来时13号元素存活的概率P(15)为，
…….依次类推
1729号元素到来时，13号元素被淘汰的概率计算公式为1729号元素被选中的概率13号元素被选中淘汰的概率，即进而可以得到1729号元素到来时13号元素存活的概率P(1729)为
根据条件概率公式，当1729号元素到来时13号元素存活的概率为：
=第13步被选中概率 不被 14号元素替换的概率 不被 15号元素替换的概率 … 不被 1729号元素替换的概率，
等价于 1 (1 - 被 14 替换的概率) (1 - 被15 替换的概率) * … (1 - 被 1729替换的概率),
等价于

根据上述计算结果可以得知，无论时10号元素之前的，还是10号元素之后的，最后被选中的概率都是相等的。因此，根据该算法可以保证数据流中的数据能够被等概率选中。

4. 代码实现：

public static class RandomBox {
    private int[] bag;  // 蓄水池具体元素值
    private int N;      // 蓄水池大小
    private int count;  // 统计每个位置被选中的次数，测试时使用该参数
    public RandomBox(int capacity) {
        bag = new int[capacity];
        N = capacity;
        count = 0;
    }
    // 等概率返回1-max中的任意数
    private int rand(int max) {
        return (int) (Math.random() * max) + 1; //
    }
    // 当前数num到达
    public void add(int num) {
        count++;
        if (count <= N) {
            bag[count - 1] = num;
        } else {
            if (rand(count) <= N) {
                // 随机数
                bag[rand(N) - 1] = num;
            }
        }
    }
    // 选择结果
    public int[] choices() {
        int[] ans = new int[N];
        for (int i = 0; i < N; i++) {
            ans[i] = bag[i];
        }
        return ans;
    }
}

测试代码：

// 100个人中随机选择十个作为幸运观众，测试50000次，统计100个观众成为候选人的次数
System.out.println("hello");
int all = 100;
int choose = 10;
int testTimes = 50000;
int[] counts = new int[all + 1];
int[] ans = new int[all + 1];
for (int i = 0; i < testTimes; i++) {
    RandomBox box = new RandomBox(choose);
    for (int num = 1; num <= all; num++) {
        box.add(num);
    }
    ans = box.choices();
    for (int j = 0; j < ans.length; j++) {
        counts[ans[j]]++;
    }
}
for (int i = 1; i < counts.length; i++) {
    System.out.println(i + " times : " + counts[i]);
}
// 幸运观众
for (int i = 0; i < ans.length; i++){
    System.out.println(ans[i]);
}

测试结果如下：
可以发现，1-100个人中每个人被选中的频率相差不大，同时返回最后选中的是为幸运观众编号。

1 times : 5106
2 times : 5032
3 times : 5055
4 times : 4963
5 times : 5022
6 times : 5080
7 times : 5036
8 times : 5103
9 times : 4987
10 times : 4926
11 times : 4960
12 times : 5073
13 times : 4984
14 times : 5020
15 times : 5124
16 times : 5066
17 times : 4993
18 times : 5035
19 times : 5110
20 times : 4907
21 times : 4996
22 times : 5067
23 times : 5002
24 times : 4964
25 times : 5008
26 times : 5097
27 times : 4884
28 times : 5082
29 times : 5019
30 times : 5022
31 times : 4890
32 times : 4898
33 times : 4925
34 times : 4901
35 times : 5038
36 times : 4893
37 times : 5109
38 times : 4963
39 times : 4930
40 times : 4886
41 times : 4910
42 times : 4985
43 times : 4911
44 times : 5040
45 times : 4997
46 times : 5026
47 times : 4943
48 times : 5165
49 times : 4858
50 times : 5031
51 times : 4930
52 times : 5019
53 times : 5074
54 times : 5062
55 times : 5015
56 times : 4953
57 times : 5049
58 times : 4925
59 times : 5053
60 times : 5133
61 times : 5036
62 times : 4935
63 times : 5059
64 times : 5005
65 times : 4973
66 times : 4921
67 times : 5026
68 times : 5035
69 times : 5019
70 times : 5019
71 times : 5064
72 times : 4917
73 times : 4972
74 times : 4955
75 times : 5046
76 times : 4874
77 times : 4989
78 times : 4958
79 times : 5112
80 times : 4864
81 times : 5039
82 times : 5093
83 times : 4968
84 times : 4933
85 times : 4931
86 times : 4964
87 times : 5096
88 times : 4940
89 times : 4935
90 times : 4926
91 times : 5017
92 times : 4978
93 times : 5042
94 times : 5061
95 times : 4928
96 times : 5084
97 times : 4976
98 times : 5033
99 times : 5027
100 times : 5015
19
35
58
94
27
36
7
13
9
96

5. 相关题目和应用：

（1）应用场景：
随机抽奖类的工程问题，如在2022年5月1日00:00:00至2022年5月3日23:59:59之间登陆的用户随机抽取100名用户作为幸运观众，发放礼品。这就可以用到蓄水池抽样算法.

（2）相关题目：

• 382. 链表随机节点
• 398. 随机数索引