水塘抽样（Reservoir sampling）

题目：给出一个数据流，这个数据流的长度很大或者未知。并且对该数据流中的数据只能访问一次。请写出一个随机选择算法，使得数据流中所有数据被选中的概率相等。

这个问题的扩展就是：如何从未知或者很大样本空间随机的取k个数？或者说，数据流长度为N行，要随机抽取k行，则每一行被抽取的概率为k/N。

这个问题也就是在大数据流中的随机抽样问题：当内存无法加载全部数据时，如何从包含未知大小的数据流中随机选取k个数据，并且要保证每个数据被抽取到的概率相等。

当 $k = 1$ 时

首先考虑最简单的情况，当 $k = 1$ 时，如何选取：

假设数据流含有N个数据，要保证每条数据被抽取到的概率相等，那么每个数被抽取的概率应该为 $\frac{1}{N}$
- 遇到第1个数 $n_{1}$ 的时候，保留它， $p (n_{1}) = 1$
- 遇到第2个数 $n_{2}$ 的时候，以 $\frac{1}{2}$ 的概率保留它，那么 $p (n_{1}) = 1 \times \frac{1}{2} = \frac{1}{2}$ ， $p (n_{2}) = \frac{1}{2}$
- 遇到第3个数 $n_{3}$ 的时候，以 $\frac{1}{3}$ 的概率保留它，那么 $p (n_{1}) = p (n_{2}) = \frac{1}{2} \times (1 - \frac{1}{3}) = \frac{1}{3} ， p (n_{3}) = \frac{1}{3}$
- …
- 遇到第i个数 $n_{i}$ 的时候，以 $\frac{1}{i}$ 的概率保留它，那么 $p (n_{1}) = p (n_{2}) = p (n_{3}) = . . . = p (n_{i - 1}) = \frac{1}{i - 1} \times (1 - \frac{1}{i}) = \frac{1}{i} ， p (n_{i}) = \frac{1}{i}$

通过以上规律可以看出，对于 $k = 1$ 的情况，数据流中第i个数被保留的概率为 $\frac{1}{i}$ 。只要采取这种策略，只需要遍历一遍数据流就可以得到采样值，并且保证所有数据被选中的概率均为 $\frac{1}{N}$ 。

当 $k > 1$ 时

对于 $k > 1$ 的情况，我们可以采取类似的策略：

假设数据流中含有N个数据，要保证每条数据被抽取到的概率相等，那么每个数被抽取的概率必然是 $\frac{k}{N}$
- 对于前k个数 $n_{1}, n_{2}, . . ., n_{k}$ ，我们保留下来，则 $p (n_{1}) = p (n_{2}) = . . . = p (n_{k}) = 1$ （下面连等采用 $p (n_{1 - k})$ 的形式
- 对于第k+1个数 $n_{k + 1}$ ，以 $\frac{k}{k + 1}$ 的概率保留它（这里只是指本次保留下来），那么前k个数中的 $n_{r} (r \in 1 - k)$ 被保留的概率可以这样表示： $p (n_{r} 被保留) = p (上一轮 n_{r} 被保留) \times (p (n_{k + 1} 被丢弃) + p (n_{k + 1} 被保留) \times p (n_{r} 未被替换))$ ，即 $p_{1 - k} = \frac{1}{k + 1} + \frac{k}{k + 1} \times \frac{k - 1}{k} = \frac{k}{k + 1}$
- 对于第k+2个数 $n_{k + 2}$ ，以 $\frac{k}{k + 2}$ 的概率保留它（这里只是指本次保留下来），那么前k+1个被保留下来的数中的 $n_{r} (r \in 1 - k + 1)$ 被保留的概率为： $p_{1 - k} = \frac{k}{k + 1} \times \frac{2}{k + 2} + \frac{k}{k + 1} \times \frac{k - 1}{k + 2}$
- …
- 对于第i(i>k)个数 $n_{i}$ ，以 $\frac{k}{i}$ 的概率保留它，前i-1个数中的 $n_{r} (r \in 1 - i - 1)$ 被保留的概率为： $p_{1 - k} = \frac{k}{i - 1} \times \frac{i - k}{i} + \frac{k}{i - 1} \times \frac{k - 1}{i} = \frac{k}{i}$

对于前k个数，全部保留，对于第i(i>k)个数，以 $\frac{k}{i}$ 的概率保留第i个数，并以 $\frac{1}{k}$ 的概率与前面已选择的k个数中的任意一个替换。

总结

也就是说，在取第i个数据的时候，生成一个0到1的随机数p，如果 $p < \frac{k}{i}$ ，替换池中任意一个为第i个数；当 $p > \frac{k}{i}$ ，继续保留前面的数。直到数据流结束，返回此k个数。但是为了保证计算准确性，一般是生成一个0到i的随机数，跟k相比。

Scala代码实现

import scala.util.Random

/**
 * @author xiaoer
 * @date 2020/1/11 23:53
 */
object ReservoirSampling {
    def main(args: Array[String]): Unit = {
        val s: Array[Int] = Array(1, 2, 3, 4, 5, 6, 7)
        val k: Array[Int] = new Array[Int](3)
        val result: Array[Int] = reservoirSampling(k, s)
        println(result.toBuffer)
    }

    /**
     * 水塘抽样算法
     *
     * @param k 抽样结果
     * @param s 样本总数
     * @return k 样本结果
     */
    def reservoirSampling(k: Array[Int], s: Array[Int]): Array[Int] = {
        // 将前 k 个数据全部抽取
        for (i <- k.indices) {
            k(i) = s(i)
        }

        // k+1 往后的数据
        for (i <- k.length until s.length) {
            val seed: Int = Random.nextInt(i)
            if (seed < k.length) {
                k(seed) = s(i)
            }
        }

        // 返回值
        k
    }
}

水塘抽样（Reservoir sampling）

水塘抽样（Reservoir sampling）

当 k = 1 k=1 k=1时

当 k > 1 k>1 k>1时

总结

Scala代码实现

当 $k = 1$ 时

当 $k > 1$ 时