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

水塘抽样

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

1. 当k=1的情况

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

P ( n i ) = 1 i ∗ ( 1 − 1 i 1 ) ∗ ( 1 − 1 i 2 ) … … ∗ ( 1 − 1 n − 1 ) ∗ ( 1 − 1 n ) = 1 i ∗ i i 1 … … ∗ n − 2 n − 1 ∗ n − 1 n = 1 n P(n_i)= \frac{1}{i} * (1-\frac{1}{i 1} )*(1-\frac{1}{i 2}) ……*(1-\frac{1}{n-1})*(1-\frac{1}{n}) \\ = \frac{1}{i} * \frac{i}{i 1} ……* \frac{n-2}{n-1} * \frac{n-1}{n} = \frac{1}{n} P(ni)=i1∗(1−i 11)∗(1−i 21)……∗(1−n−11)∗(1−n1)=i1∗i 1i……∗n−1n−2∗nn−1=n1

2. k>1的情况

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

P ( n i ) = k i ∗ [ ( 1 − k i 1 ) ( k i 1 ) ( 1 − 1 / k ) ] ∗ [ ( 1 − k i 2 ) ( k i 2 ) ( 1 − 1 / k ) ] … … ∗ [ ( 1 − k n ) ( k n ) ( 1 − 1 / k ) ] = k i ∗ i i 1 ∗ i 1 i 2 … … ∗ n − 1 n = k n P(n_i)= \frac{k}{i} * [(1-\frac{k}{i 1} ) (\frac{k}{i 1})(1-1/k) ] *[ (1-\frac{k}{i 2} ) (\frac{k}{i 2})(1-1/k) ] …… *[ (1-\frac{k}{n} ) (\frac{k}{n})(1-1/k) ]\\ = \frac{k}{i} * \frac{i}{i 1} * \frac{i 1}{i 2} …… * \frac{n-1}{n} = \frac{k}{n} P(ni)=ik∗[(1−i 1k) (i 1k)(1−1/k)]∗[(1−i 2k) (i 2k)(1−1/k)]……∗[(1−nk) (nk)(1−1/k)]=ik∗i 1i∗i 2i 1……∗nn−1=nk

代码实现

  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

这篇好文章是转载于：学新通技术网