蓄水池抽样算法/水塘采样算法

1、概念介绍

简介：

水塘抽样是一种随机抽样算法，它能够在一个很大的集合中，抽取一部分样本，并保证每个样本的选取概率都是相等并随机的。

特点：

1、选取集合可以非常大，甚至不知道边界。

2、每个样本的选取随机且概率相等。

3、时间复杂度较低，O(n)，节省内存。

适用场景：

1、在一些非常大的集合，或者未知大小的集合，不知道边界的集合，不知道文件总行数的情况下，随机抽取k个元素。

2、保证每个元素抽取都是均匀随机的并且概率相等。

3、尽量高效，节省内存地抽取

2、常见问题场景

很多大公司的面试题都考察过这个算法，下面介绍几个这样的例题

问题一：

问题：我有一个长度为N的单链表,N的值非常大,我不清楚N的确切值.我怎样能写一个尽可能高效的算法来返回链表中K个完全随机的数.
要求：
1、只能遍历一次链表
2、返回的值要均匀随机，每个的概率要相等

一般的想法就是，我先遍历一遍链表，得到链表的总长度 n，再生成一个 [1,n] 之间的随机数为索引，然后再遍历链表，找到索引对应的节点，不就是一个随机的节点了吗？

但题目说了，只能遍历一次，意味着这种思路不可行。题目还可以再泛化，给一个未知长度的序列，如何在其中随机地选择 k 个元素？想要解决这个问题，就需要著名的水塘抽样算法了。

问题二：

问题：我们有1T的文本文件存在硬盘中，想随机抽取几行，需要保证尽可能少得使用内存并且能够完全随机。

之前想到的加载到内存就不太适合了,但是还可以想到别的办法,比如每次读取一行记录加载到内存,记数+1,清空内存中行数据,直到最后统计一共多少行,然后根据总行数来计算K个随机数.如何再取回行对应的数据呢?我们可以再遍历一遍,一边遍历一边记录这一行的号码是不是在k个随机数中,如果是,则将该行内容保留.

这样的话遍历两次应该可以做到,但是1T数据遍历两次的时间消耗是非常高的.

所以还有更好的方案吗,那就是水塘抽样算法.

问题三：

问题：给定一个数据流，数据流长度N很大，且N直到处理完所有数据之前都不可知，如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出这组数据的k个概率相等的均匀抽样。
要求：
1、仅扫描数据一次
2、空间复杂度为O(N),空间复杂度与整个数据量无关，只与抽样大小有关。
3、扫描到数据的前n个数据时（n>k），保存当前已扫描数据的k个均匀抽样。

根据要求，首先体积很大内存一次装不下，不能直接不能直接取N内的k个随机数，因为N的长度是未知的。此外也不能采用不能先遍历一遍，然后分块存储数据，再随机选取。最后要求是数据选取绝对随机的保证。

可以采用水塘抽样算法：

3、水塘抽样算法原理

问题切入：

给定一个数据流，数据流长度N很大，如何从中随机选取k个数据，并且要保证每个数据被抽取到的概率相等。
注意：数据流长度N很大，内存无法加载全部数据

算法思路分析
对这个问题，如何做到均匀随机，概率相等呢？我们以什么样的方式去抽取样本呢？

首先N很大，不知道长度，所以我们无法调用random(N)来取一个1-N之间的随机数，我们只能将元素一个一个遍历加载进入内存。

比如我们已经遍历了i个元素，可以从这i个元素中随机取了一个元素a，那如果现在再给你一个新元素b，你是继续留着a呢？还是抽取b作为样本结果呢？以什么样的原则来选择a和b呢？你的选择能够保证概率相等吗？

这里就用到了我们的蓄水池抽样算法。

当k = 1时

首先考虑最简单的情况，当k = 1，只取一个元素时，如何选取：

假设数据流含有N个数据，要保证每条数据被抽取到的概率相等，那么每个数被抽取的概率应该为1/N

当遇到第1个数的时候，概率为1，直接保留第一个数为样本。

当遇到第2个数的时候，概率为1/2，第1个数和第2个数各有1/2的概率被保留。

当遇到第3个数的时候，3被保留的概率为1/3；(之前剩下的数假设1被保留)，2/3的概率 1 被保留，(此时1被保留的总概率为 2/3 * 1/2 = 1/3)

当遇到第4个数的时候，4被保留的概率为1/4，(之前剩下的数假设1被保留)，3/4的概率 1 被保留，(此时1被保留的总概率为 3/4 * 2/3 * 1/2 = 1/4)

…

当遇到第i个数的时候，i被保留的概率为1/i，之前的样本数保留不被替换的概率也为1/i

所以，依次类推，每个数被保留的概率都是1/N。

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

算法总结：

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

4、水塘抽样算法实现

算法图解：

当K=1时：
如果K=1，即每次只取一个样本值，那么逐个遍历数组，当遇到第i个数时，就以 1 / i 的概率抽取它做样本值，以（i - 1）/ i 的概率保留原来的样本值，不替换。
算法原理如下图：

当K > 1时：
如果k > 1，即要从数据流N中取K个样本值，那么要保证每条数据被抽取到的概率相等，每个数被抽取的概率必然是K / N
算法实现：对于前K个数，全部保留，作为样本集合。对于第i （i > k）个数，则以k/i 的概率抽取第i个数，并以1/k 的概率与前面已经选择的k个数中的任意一个替换。
算法原理如下图：

代码实现（java）

当K=1时：

/**
 * 当k = 1时，即只取1个数
 *
 * 算法步骤：
 * 1、遍历集合第1个元素，直接取出，作为样本数sample
 * 2、从集合第2个元素开始，使用下标 i 标记遍历的第几个元素。取一个随机数 rand = random【1，i】
 *    如果 rand == 1，则开始替换，即将当前样本元素，替换为当前遍历集合N中第i个元素。即sample = N[i]
 *    如果 rand > 1，则不替换，保留之前样本元素。
 * 3、继续往下迭代遍历，直到结束。最终样本为sample。
 */
public class ReservoirSampingTest {
    public static void main(String[] args) {
        // 定义我们要遍历的集合，原则上是无限大，我们这里简化为6个元素。
        int[] data = {1, 2, 3, 4, 5, 6};
        // 只取一个元素，所以直接定义一个变量，初始化样本值
        int sample = 0;
        // 调用抽样算法，获取样本
        sample = reservoirSamping(sample, data);
        System.out.println("最终的样本为：" + sample);
    }

    /**
     * 采样算法
     * @param sample 返回的采样的数 （k=1时，这里不传也可以）
     * @param data 遍历的大集合
     * @return 返回采样的结果
     */
    public static int reservoirSamping(int sample, int[] data) {
        // 第一个元素，直接取出
        sample = data[1];
        for (int i=1; i<data.length; i++) {
            // 定义一个随机数，范围为[1, i]，注意Random是左闭右开，所以要加1才能把i包括进去
            int rand = new Random().nextInt(i) + 1;
            // 此时随机数范围最小为1，不是0，所以当随机数等于1时，则替换为当前数
            if (rand == 1) {
                sample = data[i];
            }
        }

        // 返回采样的结果
        return sample;
    }
}

当K>1时，取k个数

跟我们上面的算法思想是一致的，只是把k=1，替换为了k=k

即概率为P = k / i

/**
* 当k > 1时，取k个数
* 算法步骤：
* 1、遍历集合前k个元素，直接取出，作为样本集合sample[]
* 2、从集合第K+1个元素开始，使用下标 i 标记遍历的第几个元素。取一个随机数 rand = random【1，i】
*    如果 rand < k，则开始替换，即将当前样本集合sample中第rand个元素，替换为当前遍历集合N中第i个元素。即 sample[rand] = N[i]
*    如果 rand > k，则不替换，保留之前样本集合sample中的元素。
* 3、继续往下迭代遍历，直到结束。最终获得k个元素的样本。
*/
public class ReservoirSampingTest {
    public static void main(String[] args) {
        // 定义我们要遍历的集合，原则上是无限大，我们这里简化为10个元素。
        int[] data = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
        // 定义我们需要抽样的样本集合，假设k=5，抽取5个元素
        int[] sample = new int[5];
        // 调用抽样算法，获取样本
        sample = reservoirSamping(sample, data);

        // 打印最终的样本集合
        for (int s : sample) {
            System.out.println("样本元素：" + s);
        }
    }

    /**
     * 水塘抽样算法
     * @param sample 要采集的样本集合
     * @param data 需要遍历大集合
     * @return 返回抽样后的样本集合sample
     */
    public static int[] reservoirSamping(int[] sample, int[] data) {
        // 假设k为样本集合的大小
        int k = sample.length;
        // 前k个元素，直接取出，放到sample中
        for (int i=0; i<k; i++) {
            sample[i] = data[i];
        }

        // 从第K+1个元素开始，根据算法随机置换sample中的元素
        for (int i=k; i<data.length; i++) {
            // 取得[1, i]之间的随机数，注意Random是左闭右开，所以要加1才能把i包括进去
            int rand = new Random().nextInt(i) + 1;
            // rand小于k，则开始置换，将sample中第rand个元素，替换为data中第i个元素
            if (rand < k) {
                sample[rand] = data[i];
            }
        }

        // 遍历一遍，采样完成，返回样本集合sample
        return sample;
    }
}

5、LeetCode中算法案例

【382. 链表随机节点】：https://leetcode-cn.com/problems/linked-list-random-node/

题解查看：https://www.cnblogs.com/doublexi/p/15666056.html

【398. 随机数索引】：https://leetcode-cn.com/problems/random-pick-index/

题解查看：https://www.cnblogs.com/doublexi/p/15666221.html

6、spark分区器（水塘采样应用案例）

最近在学习spark的分区器的时候，发现一个很有意思的东西水塘采样算法，又称蓄水池抽样算法。

spark的分区器主要有三种：

（只有key-value类型的RDD才有分区器，Value类型的RDD分区器的值是None。）

1、HashPartitioner：数据分区规则为 partitionId = Key.hashCode % numPartitions。

特点：采用哈希的方式将同一类型的Key分配到同一个Partition中。

但是当某个key特别多时，容易导致某个分区数非常多，导致各个分区数据不均匀，形成数据倾斜。

2、RangePartitioner：将一定范围内的数映射到某一个分区内。分区与分区之间数据是有序的，但分区内的元素是不能保证顺序的。sortByKey会使用RangePartitioner。

特点：能够尽量保证每个分区中数据量的均匀。（主要用于数值类型）

3、自定义分区器。

你可以通过partitionBy来指定自己想要用哪种分区器。

关于RangePartitioner：

现在的问题：

在RangePartitioner中，在执行分区之前其实并不知道数据的分布情况，不知道数据的总量，不知道数据的范围边界（这个边界很重要），也不知道数据分布在哪个范围。

那我们就很难给数据分区，把哪些数据，或者是哪一段范围的数据分配到对应分区，从而保证各个分区的数据量都比较均匀。因为不知道数据的分布范围，不好给数据划分边界，也不好分区。

这时，就可以用到水塘采样算法了。

RangePartitioner在分区的时候，不知道数据总量，数据分布情况，但是它需要知道这些去划分分区后的边界。所以它会调用水塘采样算法，返回RDD的总采样数据量，分区ID和，每个分区数据总量，每个分区的采样数据。

RangePartitioner中是对水塘采样算法，进行了变种和加强，有兴趣的可以参考：

https://www.cnblogs.com/tongxupeng/p/10435976.html

它的主要步骤为：

1、计算总体的数据抽样大小sampleSize，计算规则是：至少每个分区抽取20个数据或者最多1M的数据量。

2、根据sampleSize和分区数量计算每个分区的数据抽样样本数量最大值sampleSizePrePartition

3、根据以上两个值进行水塘抽样，返回RDD的总采样数据量，分区ID和，每个分区数据总量，每个分区的采样数据。

4、数据抽样完成后，需要对不均衡的Partition重新进行抽样，默认当Partition中包含的数据量大于平均值的三倍时，该Partition是不均衡的。当采样完成后，利用样本容量和RDD中包含的数据总量，可以得到整体的一个数据采样率fraction。利用此采样率对不均衡的Partition调用sample算子重新进行抽样。(此时，各个分区的数据规模已经知道了，就没必要再水塘采样了，直接调用sample算子二次采样)

5、根据各个分区的key与权重值，调用determineBounds方法计算出步长与边界

6、计算每个Key所在Partition：当分区范围长度在128以内，使用顺序搜索来确定Key所在的Partition，否则使用二分查找算法来确定Key所在的Partition。