水塘抽样算法及其代码实现(Scala)

水塘抽样

水塘抽样是一系列的随机算法，其目的在于从包含n个项目的集合S中选取k个样本，其中n为一很大或未知的数量，尤其适用于不能把所有数据都存放到内存的情况，或者流式数据中的随机抽样

1. 当k=1的情况

要保证每个数等概率被抽中，那么每个数抽中的概率为1/N ，抽样策略为：以1/i 的概率保留第i个数：

$$\begin{gathered} P(n_i)=\frac{1}{i}\ast (1-\frac{1}{i+1})\ast (1-\frac{1}{i+2})\dots \dots \ast (1-\frac{1}{n-1})\ast (1-\frac{1}{n}) \\ =\frac{1}{i}\ast \frac{i}{i+1}\dots \dots \ast \frac{n-2}{n-1}\ast \frac{n-1}{n}=\frac{1}{n} \end{gathered}$$

2. k>1的情况

k>1的情况要保证每个数倍抽中的概率都为 k/n ,此时的抽样策略为：第i个数我们以 k/i 的概率保留下来，并且随机替换前k个选中的数的一个。

$$\begin{gathered} P(n_i)=\frac{k}{i}\ast [(1-\frac{k}{i+1})+(\frac{k}{i+1})(1-1/k)]\ast [(1-\frac{k}{i+2})+(\frac{k}{i+2})(1-1/k)]\dots \dots \ast [(1-\frac{k}{n})+(\frac{k}{n})(1-1/k)] \\ =\frac{k}{i}\ast \frac{i}{i+1}\ast \frac{i+1}{i+2}\dots \dots \ast \frac{n-1}{n}=\frac{k}{n} \end{gathered}$$

代码实现

  def reservoirSample(arr: Array[Int], k: Int = 1) = {

    // 前k个数，首先都会保留下来
    val res: Array[Int] = arr.take(k)
    for (i <- k until arr.length) {
      val c = math.random
      //  第 i 个数 （i>k） 以 k/i 的概率保留下来;
      if ((k / (i + 1.0)) >= c) {
        //  随机替换前面已选择的k个数中的任意一个 (每个数被替换的概率为1/k)
        res(Random.nextInt(k)) = arr(i)
      }
    }
    res
  }
  

数据验证：

  
  def main(args: Array[String]): Unit = {

    def fibFrom(a: Int, b: Int): Stream[Int] = a #:: fibFrom(b, a + b)

    val fbs = fibFrom(1, 2).take(8).toArray



    val fb = (1 to 80000).flatMap(x => reservoirSample(fbs,3).map(y => (y, 1)))
      .groupBy(_._1)
      .map(t => (t._1, t._2.map(_._2).sum))
      .toList.sortBy(_._1)

    println(fb)

 

数据结果：从数据流中抽取的数据结果符合预期

List((1,30132), (2,29966), (3,30100), (5,29941), (8,30089), (13,29827), (21,30127), (34,29818))

参考：

https://zhuanlan.zhihu.com/p/29178293?utm_source=wechat_timeline