蓄水池抽样：动态等概率抽样

前言

昨天学习了一个比较有意思的也比较简单的算法，记录一下

解决的问题

不知道总样本为多少, 而且由于内存不足，不能把样本全部读进来；
现在想一个办法，从所有样本中随机等概率获取其中一个样本输出
【或者也可以：对于当前已经遍历过的样本，从中等概率的得到一个样本】

表达1实际上是表达2的一个特殊情况【即当前固定为最后一轮】

算法

经典算法：将所有的遍历样本存入到一个集合中，每次遍历对集合元素随机取值。

这显然是不符合题意的：对于要求最后一轮随机取样，就需要将所有样本都放入集合，而内存是不足的，因此无法做到
即使只在内存放下标，若样本数量足够多的情况下【假设达到了亿的级别】，只是元素下标全部放入内存都会把内存撑爆。

我们需要考虑一种空间复杂度足够低的方法，一种有效的方法就是蓄水池抽样。

算法描述

设定两个量：1）res，存储当前抽样结果；2）i：当前遍历的样本数量【从1开始】

1. 逐一遍历元素，设当前遍历到的元素为i；
2. 取一个[0, 1)之间随机数r，若$r<\frac{1}{i}$, 将res替换为n；
3. 循环结束后，res即为结果

对于要求1，我们可以在每轮替换后输出一次，就相当于获取了每次遍历的等概率抽样。
很简单的算法，证明过程也很有意思。

证明

数学归纳法：

1. 设第i轮时，当前遍历了一共i个元素,抽取其中每个元素的概率都是等概率的，即$\frac{1}{i}$
2. 则，需要证明第$i+1$轮时，$i+1$个元素每个被抽取的概率都是$\frac{1}{i+1}$
3. 首先证明，抽取第$i+1$个物品的概率为$\frac{1}{i+1}$：
   1. 我们抽取第$i+1$物品，是通过一个随机数$r=random([0,1))$，若$r<\frac{1}{i+1}$, 则进行替换, 或者说是抽取；
   2. 亦即证明$r<\frac{1}{i+1}$的概率为$\frac{1}{i+1}$
   3. 我们知道random函数是一个均匀分布，而$r<\frac{1}{i+1}$就是说r占据了[0,1)的$[0,\frac{1}{i+1})$区域，根据均匀分布定义，其区间的占比就是概率，而$[0,\frac{1}{i+1})$占据[0, 1)区间的$\frac{1}{i+1}$，即$r<\frac{1}{i+1}$概率为$\frac{1}{i+1}$
4. 接下来证明其他i个元素选中的概率也是$\frac{1}{i+1}$：
   1. 其他元素（不失一般性，设为j）要想被选中，要满足两个条件：1）之前那轮res为j；2）这一轮没有被i+1替换；
   2. 1）的概率为$\frac{1}{i}$【题设】；2）的概率为$1-\frac{1}{i+1}=\frac{i}{i+1}$, 这两件事是分步进行的，符合乘法原则，最终概率为$\frac{1}{i}\times \frac{i}{i+1}=\frac{1}{i+1}$
5. 综合3,4，得证。

伪代码

    public int sample(int[] nums) {
        int res = -1;
        int i = 0;
        for (int n : nums) {
            i++;
            double r = Math.random();
            if (r < 1 / (double) i) {
                res = n;
            }
        }
        return res;
    }

扩展问题1

从未知数量的一批样本中，等概率的取出k个物品，必须保证每个物品选中的概率都是相等的(根据概率公式，从N个物品中取出k个，每个物品被选中的概率为$\frac{k}{N}$) 【为了简化问题，我们总是假设$N>k$】

等价问题：
遍历这N个元素：
1）若当前遍历的样本数量小于等于k个，则全部选取；
2）若当前遍历的样本数量大于k个，选取其中的k个，每个选取的概率都是独立且相等的(全为i/k)。

算法

设定两个量：1）res = []，存储当前抽样结果集合；2）i：当前遍历的样本数量【从1开始】

1. 逐一遍历元素，设当前遍历到的元素为i；
2. 若res中结果不足k个，直接加入res;
3. 若结果超过k个，取一个[0, 1)之间随机数r，若$r<\frac{k}{i}$, 决定将res中某个元素替换为i；
4. 替换也是抽取一个随机数，随机为[1, k]之间的整数，即res中的下标；
5. 循环结束后，res即为结果

证明

仍然采用数学归纳法：【只讨论$i>k$的情况，因为之前都是必选的，概率都是1】

1. 设第i轮时，当前遍历了一共i个元素,抽取其中每个元素的概率都是等概率的，即$\frac{k}{i}$
2. 则，需要证明第$i+1$轮时，$i+1$个元素每个被抽取的概率都是$\frac{k}{i+1}$
3. 首先证明，抽取第$i+1$个物品的概率为$\frac{k}{i+1}$：
   1. 这个和前面的均匀分布证明是类似的，不再赘述。
4. 接下来证明其他i个元素选中的概率也是$\frac{1}{i+1}$：
   1. 其他元素（不失一般性，设为j）要想被选中，要满足两个条件：1）之前那轮j在res中；2）这一轮i+1没有拿到替换的资格(设为II条件)或者虽然i+1拿到了替换资格，但是替换的不是j(条件III)；
   2. 1）的概率为$\frac{k}{i}$【题设】；2）的概率分为两个部分，II条件的概率是$1-\frac{k}{i+1}=\frac{i+1-k}{i+1}$, 条件III概率是$\frac{k-1}{k}\times \frac{k}{i+1}=\frac{k-1}{i+1}$。条件II和条件III因为是或事件，符合加法原则，因此2）的概率为$\frac{k-1}{i+1}+\frac{i+1-k}{i+1}=\frac{i}{i+1}$。1）2）这两件事是分步进行的，符合乘法原则，最终概率为$\frac{k}{i}\times \frac{i}{i+1}=\frac{k}{i+1}$
5. 综合3,4，得证。

伪代码

    public int[] sample2(int[] nums, int k) {
        int[] res = new int[k];
        int i = 0;
        for (int n : nums) {
            if (i < k) {
                res[i++] = n;
            } else {
                i++;
                if (Math.random() < k / (double)i) {
                    res[(int)(Math.random() * k)] = n;
                }
            }
        }
        return res;
    }

优化及证明

我们上面的代码中，将替换的步骤拆分为了两步，做了两次随机，实际上，可以只做一次随机，让第二部随机也能直接使用第一步随机的结果，但是第二步随机要求是整数，而第一步随机是[0,1)之间的小数啊？很简单，等量放大即可

算法2中采用的r的抽样和判定的条件为：
$r<\frac{k}{i}$，两边同时乘以i，即$r\times i<k$
由于r是一个[0,1)之间的随机数，因此$r\times i$就是[0, i)之间随机数，右边的上界也被放大了，整体不等式是与之前的等价的；
而若$r\times i<k$，$r\times i$就是解集res合法的下标，即要替换的位置。

伪代码

    public int[] sample3(int[] nums, int k) {
        int[] res = new int[k];
        int i = 0;
        for (int n : nums) {
            if (i < k) {
                res[i++] = n;
            } else {
                i++;
                int r = (int)(Math.random() * i);
                if (r < k) {
                    res[r] = n;
                }
            }
        }
        return res;
    }

例题

lc 398

class Solution {
    int[] nums;
    public Solution(int[] nums) {
        this.nums = nums;
    }
    
    public int pick(int target) {
        int i = 0, res = -1, cnt = 0;
        for (int n : nums) {
            if (n == target) {
                i++;
                if (Math.random() < 1 / (double) i) {
                    res = cnt;
                }
            }
            cnt++;
        }
        return res;
    }
}

但是目前这个O(n^2)的算法会超时。

lc 382

class Solution {
    ListNode head;
    public Solution(ListNode _head) {
        head = _head;
    }
    
    public int getRandom() {
        int i = 0, res = 0;
        ListNode cur = head;
        while (cur != null) {
            i++;
            if (Math.random() < 1 / (double)i) {
                res = cur.val;
            }
            cur = cur.next;
        }
        return res;
    }
}

lc497

class Solution {

    int[][] rects;
    public Solution(int[][] rects) {
        this.rects = rects;
    }
    
    public int[] pick() {
        double sum = 0; //当前遍历的总点数
        int[] res = null;
        for (int[] r : rects) {
            long cur = (r[2] - r[0] + 1) * (r[3] - r[1] + 1);
            sum += cur;
            if (Math.random() < cur / sum) {
                res = r;
            }
        }
        //从res中随机取出一点
        int w = res[2] - res[0] + 1, h = res[3] - res[1] + 1;
        return new int[] {res[0] + (int)(Math.random() * w), res[1] + (int)(Math.random() * h)};
    }
}

这题改变挺大的，有必要解释一下。
最大的不同点在于，抽取的目标是一个矩形，但是每个矩形的大小是不等的，在第i轮获取矩形j(j属于[1, i])的概率是不等的，而是一个加权的概率，权重为其含有的点的数量。
但是我们只需要考虑需不需要用第i个矩形来替代，所以只需要计算第i个矩形的概率权重即可。

总结

蓄水池抽样不是一种解决概率问题很好的方法，只是一种以时间换空间的方法【往往时间复杂度会上升一个层次】，这更多只会出现在面试的问答之中，看你掌握的算法的广度，真正笔试还是想别的办法吧【比如第一题就可以用一个Map缓存，第二题可以用一个List缓存，第三题可以用前缀数组加二分，效果都比这个算法好】。