【数据结构与算法】蓄水池抽样算法（Reservoir Sampling）
问题描述
给定一个数据流，数据流长度 N 很大，且 N 直到处理完所有数据之前都不可知，请问如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出 m 个不重复的数据。
比较直接的想法是利用随机数算法，求 random(N) 得到随机数，但是题目表明数据流极大，这种大数据量是无法一次都读到内存的，这就意味着不能像数组一样根据索引获取元素。获取 N 只能对所有数据进行遍历，耗费时间较大，并且题目强调只能遍历一遍，意味着不能先获取到 N ，那么采用分块存储数据的方法也不可取（遍历不止一遍）；如果采用估算，可能导致采样数据不平均。
蓄水池抽样算法
假设数据序列的规模为 n（蓄水池大小），需要采样的数量的为 k。
首先构建一个可容纳 k 个元素的数组，将序列的前 k 个元素放入数组中。
然后从第 k+1 个元素开始，以 k/n 的概率（n 为当前索引位置）来决定该元素是否被替换到数组中（数组中的元素被替换的概率是相同的）。当遍历完所有元素之后，数组中剩下的元素即为所需采取的样本。
证明
- 对于第
i个元素（i <= k)，该元素被选中的概率为1，且索引idx <= k时，该元素被替换的概率为0，当索引idx走到k + 1时，第k + 1个元素被选中进行替换的概率为 $ \frac{\mathrm{k}}{\mathrm{k}+1}$ ，第i个元素被选中的概率为 $ \frac{\mathrm{1}}{\mathrm{k}}$，于是第i个元素被第k + 1个元素替换的概率为 $ \frac{\mathrm{k}}{\mathrm{k}+1}$ * $ \frac{\mathrm{1}}{\mathrm{k}}$ = $ \frac{\mathrm{1}}{\mathrm{k}+1}$ ，则第i个元素不被替换的概率为 1 - $ \frac{\mathrm{1}}{\mathrm{k}+1}$ = $ \frac{\mathrm{k}}{\mathrm{k}+1}$ ，同理，第k + 2个元素被选中进行替换的概率为 $ \frac{\mathrm{k}}{\mathrm{k}+2}$ ， 第i个元素被选中的概率为 $ \frac{\mathrm{1}}{\mathrm{k}}$ ，于是第i个元素被第k + 2个元素替换的概率为 $ \frac{\mathrm{k}}{\mathrm{k}+2}$ * $ \frac{\mathrm{1}}{\mathrm{k}}$ = $ \frac{\mathrm{1}}{\mathrm{k}+2}$ ，第i个元素不被替换的概率为 1 - $ \frac{\mathrm{1}}{\mathrm{k}+2}$ = $ \frac{\mathrm{k}+1}{\mathrm{k}+2} $。以此类推，运行到第n步时，被保留的概率 = 被选中的概率 * 不被替换的概率，即：
\[1\times \frac{\mathrm{k}}{\mathrm{k}+1}\times \frac{\mathrm{k}+1}{\mathrm{k}+2}\times \frac{\mathrm{k}+2}{\mathrm{k}+3}\times ...\times \frac{\mathrm{n}-1}{\mathrm{n}}=\frac{\mathrm{k}}{\mathrm{n}}
\]
- 对于第
j个元素（j > k），该元素被选中的概率为 $ \frac{\mathrm{k}}{\mathrm{j}}$ ，第j + 1个元素被选中进行替换的概率为 $ \frac{\mathrm{k}}{\mathrm{j}+1}$，第j个元素被选中的概率为 $ \frac{\mathrm{1}}{\mathrm{k}}$，第j个元素被替换的概率为 $ \frac{\mathrm{k}}{\mathrm{j}+1}$ * $ \frac{\mathrm{1}}{\mathrm{k}}$ = $ \frac{\mathrm{1}}{\mathrm{j}+1}$，则第 j 个元素不被替换的概率为 1 - $ \frac{\mathrm{1}}{\mathrm{j}+1}$ = $ \frac{\mathrm{j}}{\mathrm{j}+1}$。则运行到第n步时，被保留的概率 = 被选中的概率 * 不被替换的概率，即：
\[\frac{\mathrm{k}}{\mathrm{j}}\times \frac{\mathrm{j}}{\mathrm{j}+1}\times \frac{\mathrm{j}+1}{\mathrm{j}+2}\times \frac{\mathrm{j}+2}{\mathrm{j}+3}\times ...\times \frac{\mathrm{n}-1}{\mathrm{n}}=\frac{\mathrm{k}}{\mathrm{n}}
\]
所以对于其中每个元素，被保留的概率都为 $ \frac{\mathrm{k}}{\mathrm{n}}$
代码实现
public class ReservoirSampling {
private int[] pool; // 蓄水池，包含所有数据
private int size; // 蓄水池规格
private Random random;
public ReservoirSampling(int size) {
this.size = size;
random = new Random();
// 初始化数据
pool = new int[size];
for (int i = 0; i < size; i++) {
pool[i] = i;
}
}
public int[] sampling(int K) {
int[] result = new int[K];
for (int i = 0; i < K; i++) { // 前 K 个元素直接放入数组中
result[i] = pool[i];
}
for (int i = K; i < size; i++) { // K + 1 个元素开始进行概率采样
int r = random.nextInt(i + 1); // 索引下标为 i 个数据时第 i + 1 个数据，r = [0,i]
if (r < K) { // 选中概率为 k/i+1
result[r] = pool[i];
}
}
return result;
}
}
测试
public static void main(String[] args) {
ReservoirSampling test = new ReservoirSampling(1000);
int[] sampling = test.sampling(5);
for (int i : sampling) {
System.out.print(i + " ");
}
}
// 输出 205 907 986 696 443，每次运行结果不同
题目
LeetCode 382. 链表随机节点
给你一个单链表，随机选择链表的一个节点，并返回相应的节点值。每个节点 被选中的概率一样 。
实现 Solution 类：
Solution(ListNode head) 使用整数数组初始化对象。
int getRandom() 从链表中随机选择一个节点并返回该节点的值。链表中所有节点被选中的概率相等。
示例：
输入
["Solution", "getRandom", "getRandom", "getRandom", "getRandom", "getRandom"]
[[[1, 2, 3]], [], [], [], [], []]
输出
[null, 1, 3, 2, 2, 3]
解释
Solution solution = new Solution([1, 2, 3]);
solution.getRandom(); // 返回 1
solution.getRandom(); // 返回 3
solution.getRandom(); // 返回 2
solution.getRandom(); // 返回 2
solution.getRandom(); // 返回 3
// getRandom() 方法应随机返回 1、2、3中的一个，每个元素被返回的概率相等。
提示：
链表中的节点数在范围 [1, 104] 内
-104 <= Node.val <= 104
至多调用 getRandom 方法 104 次
进阶：
如果链表非常大且长度未知，该怎么处理？
你能否在不使用额外空间的情况下解决此问题？
解：典型的蓄水池算法，当 k 取 1 时的特殊情况，每次只取出一个元素。
class Solution {
ListNode head;
Random random = new Random();
public Solution(ListNode _head) {
this.head = _head;
}
// 另第 idx 个结点被选中的概率为 1/idx ，则该结点不被后面结点覆盖的概率为 1/idx *
// (1 - 1/(idx+1)) * (1 - 1/(idx+2)) * ...* (1 - 1/n) = 1/n
// 白话：对第 idx 个结点计算概率，random = [0, idx)， 则random = 0 的概率为 1/idx
// 只要第 idx 个结点的 random 为 0 则选中覆盖原答案，直到选到最后一个结点。
public int getRandom() {
int idx = 1;
ListNode node = head;
int ans = node.val;
while(node != null) {
if(random.nextInt(idx) == 0) ans = node.val;
node = node.next;
idx++;
}
return ans;
}
}
LeetCode 398. 随机数索引
给定一个可能含有重复元素的整数数组，要求随机输出给定的数字的索引。 您可以假设给定的数字一定存在于数组中。
注意：
数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。
示例:
int[] nums = new int[] {1,2,3,3,3};
Solution solution = new Solution(nums);
// pick(3) 应该返回索引 2,3 或者 4。每个索引的返回概率应该相等。
solution.pick(3);
// pick(1) 应该返回 0。因为只有nums[0]等于1。
solution.pick(1);
解：只需要考虑给定数字即可，对遍历到的给定数字进行编号（1，2，...），再按照蓄水池算法随机取出一个即可
class Solution {
private int[] nums;
public Solution(int[] nums) {
this.nums = nums;
}
public int pick(int target) {
int ans = 0;
int idx = 0;
Random random = new Random();
for(int i = 0; i < nums.length; i++) {
if(nums[i] == target) {
idx++;
if(random.nextInt(idx) == 0) ans = i;
}
}
return ans;
}
}
