水塘抽样 Reservoir sampling

什么是水塘抽样？

水塘抽样是从n个元素中随机选取k个元素的算法。其中n可以是一个非常大的或者未知的数字。通常来说，水塘抽样算法用于n超过内存的容量或n是一个非常大的输入流的抽样。

从已知数目的n个元素中抽样k个

算法步骤：
1. 首先将n个元素中的前k个放入存储返回值的数组（记为ret)中。
2. 对于从剩余的n - k个元素，从第k+1个元素开始，记录当前已考虑过的元素的count，然后从0到count - 1之间取随机数记为rand，如果rand小于k，则将ret[rand - 1]赋值为当前元素。
3. 直到将n个元素全部循环一遍。

Java实现：

public class ReservoirSample {
    public int[] reservoirSample(int[] input, int k) {
        int n = input.length;
        int[] ret = new int[k];
        for (int  i = 0; i < n; i++) {
            if (i < k) {
                ret[i] = input[i];
            } else {
                int rand = new Random().nextInt(i);
                if (rand < k) {
                    ret[rand] = input[i];
                }
            }
        }

        return ret;
    }
}

从未知长度的输入流中抽样k个元素

如果我们需要抽样的输入为一个流输入，那么我们就无法事先知道这个流的长度。但考虑我们上面的算法，就会发现其实其也适用于流输入。

相似地，我们任然用一个count来记录当前遇到的元素的总数。当count小于k的时候，我们就直接将当前元素放入ret，否则我们仍取0到count-1之间的随机数rand，当rand小于k的时候，我们将其放入ret[rand]。

证明

对于k=1的情况，很容易通过数学归纳法证明。

基本情况:
当仅有一个元素被接收到时，其被选择的概率为：1/1 = 1.
递推:
假设有i个元素就被收到，每一个元素被选中的概率为1/i.
对于第(i+1)个元素而言, 当且仅当rand(i+1) == 0这个元素才会被选中，概率为1/(i + 1).
那么当接收到i+1个元素时，第i个元素被选中的概率为 (1/i) * i/(i + 1) = 1/(i + 1).。

因此可得当共有n个元素被收到时，每一个元素被选中的概率都为 1/n。

我们也可以这样理解这个问题：
当共有n个元素出现，如果我们选中了第i个元素，对于所有从第 (i+1)到第n个元素，rand(j)必须不为0，这一概率为P=1i×ii+1×i+1i+2×⋯×n−1n=1