【随机算法】加权蓄水池抽样

摘要: 加权蓄水池抽样

【对数据分析、人工智能、金融科技、风控服务感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：潮汐朝夕
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

---

各位好，在文章 【随机算法】蒙特卡洛 中，我们了解了蒙特卡洛方法的基本思想，实现蒙特卡洛方法最重要的是根据需求正确地实现对给定分布的随机变量的采样，并学习了直接采样。

在 【随机算法】拒绝采样 中我们学习了拒绝采样，在 【随机算法】蓄水池抽样 中我们学习了蓄水池抽样。

本文我们来看加权蓄水池抽样。我们首先给出加权蓄水池抽样的适用场景、算法流程，以及正确性证明。然后用加权蓄水池抽样解决力扣上的一道随机算法的题目，528，此外本题也可以用二分解决。

---

加权蓄水池抽样

场景

给定一个数据流 A，数据流长度 N 很大，且 N 直到处理完所有数据之前都不可知。

每个数据流中的每个数据 A[i] 有一个权重 w[i]，在处理当前数据的时候，后面数据的权重未知。

如何在只遍历一遍数据 $O(N)$ 的情况下，能够随机选取出 m 个不重复的数据。

• 数据流长度未知，可能很大，将全部数据读入内存的算法不可行。
• 所有数据过一遍之后出解，时间复杂度为 $O(N)$。
• 随机选取 m 个数，每个数被选中的概率与其权重成正比。

算法

输入：数据流 A，长度未知，第 i 个样本 Ai 的权重为 wi
输出：蓄水池 R 长度为 m

枚举数据流 Ai in A :
   ri = rand(0, 1) ^ (1 / wi)
   if i <= m:
      (Ai, ri) 加入 R
   else:
      (At, rt) 为 R 中 r 最小的那个样本
      if ri > rt:
         (Ai, ri) 替换 (At, rt)

• 当 wi 值为一般权重而非概率值时，可能会是一个很大的数值，使得求 ri 的指数操作丢失精度。

解法： ri 取 log() 而变成 ri = log(rand(0, 1)) / wi，因为后续在各个 ri 之间只涉及比较相对大小而不是绝对值，所以可以保证精度的同时不影响结果。

• 用小顶堆维护蓄水池中持有的随机数最小的哪个样本的查询，基于小顶堆的算法的时间复杂度就是 $O(m\log N)$

正确性证明

算法的核心是 $ri = rand(0, 1)^{1/wi}$

• 令 U1, U2 为两个相互独立的随机变量且均服从 [0, 1] 区间上的均匀分布
• 令 $r1 = (U1)^{1/w1}$, $r2 = (U2)^{1/w2}$, 其中样本权重 w1, w2 > 0
• 可以推出概率关系：$P(r1 \leq r2) = \frac{w2}{w1+w2}$

更多的推导参考 2006 年提出时的论文原文 《Weighted random sampling with a reservoir》

Follow Up

下面两个 Follow Up 可以思考一下。

• 分布式 A-res。
• 有放回的情况: 用 m 个长为 1 的独立蓄水池。

下面我们看例题，解决方法就是蓄水池抽样，算法流程与前面介绍的一样，因此后面的题目直接给出代码。

题目

• 528. 按权重随机选择

给你一个 下标从 0 开始 的正整数数组 w ，其中 w[i] 代表第 i 个下标的权重。

请你实现一个函数 pickIndex ，它可以 随机地 从范围 [0, w.length - 1] 内（含 0 和 w.length - 1）选出并返回一个下标。选取下标 i 的 概率 为 w[i] / sum(w) 。

例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25 （即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。

提示：

1 <= w.length <= 1e4
1 <= w[i] <= 1e5
pickIndex 将被调用不超过 1e4 次

示例 1：
输入：
[“Solution”,”pickIndex”]
[[[1]],[]]
输出：
[null,0]
解释：
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。

示例 2：
输入：
[“Solution”,”pickIndex”,”pickIndex”,”pickIndex”,”pickIndex”,”pickIndex”]
[[[1,3]],[],[],[],[],[]]
输出：
[null,1,1,1,1,0]
解释：
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // 返回 1，返回下标 1，返回该下标概率为 3/4 。
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 0，返回下标 0，返回该下标概率为 1/4 。
由于这是一个随机问题，允许多个答案，因此下列输出都可以被认为是正确的:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
……
诸若此类。

算法1：二分

对加权随机事件的二分也是随机算法中比较常见的一种算法，本题就是常规的已知左右样本和权重的二分法加权采样。算法细节参考：

【随机算法】加权随机事件的二分查找。

代码如下：

class Solution {
public:
    Solution(vector<int>& w) {
        int sum = 0;
        for(int i: w) sum += i;
        int n = w.size();
        ranges.assign(n, 0.0);
        int pre_sum = w[0];
        for(int i = 1; i < n; ++i)
        {
            ranges[i] = (double)pre_sum / sum;
            pre_sum += w[i];
        }
        dre = std::default_random_engine();
        dr = std::uniform_real_distribution<double>(0.0, 1.0);
    }

    int pickIndex() {
        double p = dr(dre);
        auto it = --upper_bound(ranges.begin(), ranges.end(), p);
        return it - ranges.begin();
    }

private:
    vector<double> ranges;
    std::default_random_engine dre;
    std::uniform_real_distribution<double> dr;
};

算法2: 加权蓄水池抽样

class Solution {
public:
    Solution(vector<int>& w) {
        this -> A = w;
        dre = std::default_random_engine();
        dr = std::uniform_real_distribution<double>(0.0, 1.0);
    }

    int pickIndex() {
        int n = (this -> A).size();
        int m = 1;
        int At = -1; // 蓄水池中的样本
        double rt = -1; // 蓄水池中样本 At 只有的随机数
        for(int i = 0; i < n; ++i)
        {
            int Ai = i;
            int wi = (this -> A)[i];
            double ri = log(dr(dre)) / (double)wi;
            if(i < m)
            {
                At = Ai;
                rt = ri;
            }
            else
            {
                if(rt < ri)
                {
                    At = Ai;
                    rt = ri;
                }
            }
        }
        return At;
    }

private:
    vector<int> A;
    std::default_random_engine dre;
    std::uniform_real_distribution<double> dr;
};

---