spark抽样之蓄水池抽样

spark随机森林的实现中用到了蓄水池抽样，主要是用在特征集抽样中，本文对蓄水池抽样的原理和spark源码进行简单分析。
1. 蓄水池抽样
形象的说，所谓蓄水池抽样就像是在向池中加水的过程中进行抽样，类比于对数据流进行抽样。前提是我们知道要抽取的样本量，但是很难知道样本空间的规模，比如流式数据，或者样本空间较大，遍历一遍代价较高，因此无法计算其抽取概率或比例，不能用常规的方法抽取数据。蓄水池抽样能在o(n)的复杂度对数据进行等概率抽样，其伪代码描述如下

sample size: kfor i = 0 to N    if i < k        choose ith value    else        M = random(0,i)        if M < k            swap Mth value and ith valueend for

2.等概性证明
在蓄水池抽样中，每个样本被抽取的概率为k/N，从伪代码中我们能看到，样本被抽取的概率可以分成两部分，一部分是样本被抽中的概率p1，另外是样本不被其后的样本替换的概率p2，样本被抽取的概率是p1*p2。
当样本index小于抽取量k时（0，k-1），p1=1，第i个样本不被第k个样本替代的概率是第k个元素产生的随机数不是该样本的index i，其概率为k/(k+1)，同理不被第k+1个样本替代的概率为(k+1)/(k+2)，因此

其抽样概率为k/N
当index大于等于k时(k, N)，第i个样本被选中的概率为

不被其后元素替代的概率

因此被抽取的概率同样为k/N
3. spark中蓄水池抽样的实现
实现的源码在org.apache.spark.util.random.SamplingUtils中的reservoirSampleAndCount函数，代码比较简单

/**   * Reservoir sampling implementation that also returns the input size.   *   * @param input input size   * @param k reservoir size   * @param seed random seed   * @return (samples, input size)   */  def reservoirSampleAndCount[T: ClassTag](      input: Iterator[T],      k: Int,      seed: Long = Random.nextLong())    : (Array[T], Long) = {    val reservoir = new Array[T](k)    // Put the first k elements in the reservoir.    var i = 0    while (i < k && input.hasNext) {      val item = input.next()      reservoir(i) = item      i += 1    }    // If we have consumed all the elements, return them. Otherwise do the replacement.    if (i < k) {      // If input size < k, trim the array to return only an array of input size.      val trimReservoir = new Array[T](i)      System.arraycopy(reservoir, 0, trimReservoir, 0, i)      (trimReservoir, i)    } else {      // If input size > k, continue the sampling process.      var l = i.toLong      val rand = new XORShiftRandom(seed)      while (input.hasNext) {        val item = input.next()        val replacementIndex = (rand.nextDouble() * l).toLong        if (replacementIndex < k) {          reservoir(replacementIndex.toInt) = item        }        l += 1      }      (reservoir, l)    }  }