【数据结构与算法】蓄水池抽样算法（Reservoir Sampling）

2022-06-05 04:09:34

问题描述

给定一个数据流，数据流长度 N 很大，且 N 直到处理完所有数据之前都不可知，请问如何在只遍历一遍数据（O(N)）的情况下，能够随机选取出 m 个不重复的数据。

比较直接的想法是利用随机数算法，求 random(N) 得到随机数，但是题目表明数据流极大，这种大数据量是无法一次都读到内存的，这就意味着不能像数组一样根据索引获取元素。获取 N 只能对所有数据进行遍历，耗费时间较大，并且题目强调只能遍历一遍，意味着不能先获取到 N ，那么采用分块存储数据的方法也不可取（遍历不止一遍）；如果采用估算，可能导致采样数据不平均。

蓄水池抽样算法

假设数据序列的规模为 n（蓄水池大小），需要采样的数量的为 k。

首先构建一个可容纳 k 个元素的数组，将序列的前 k 个元素放入数组中。

然后从第 k+1 个元素开始，以 k/n 的概率（n 为当前索引位置）来决定该元素是否被替换到数组中（数组中的元素被替换的概率是相同的）。当遍历完所有元素之后，数组中剩下的元素即为所需采取的样本。

证明

对于第 i 个元素（i <= k)，该元素被选中的概率为 1，且索引 idx <= k 时，该元素被替换的概率为 0 ，当索引 idx 走到 k + 1 时，第 k + 1 个元素被选中进行替换的概率为 $ \frac{\mathrm{k}}{\mathrm{k}+1}$ ，第 i 个元素被选中的概率为 $ \frac{\mathrm{1}}{\mathrm{k}}$，于是第 i 个元素被第 k + 1 个元素替换的概率为 $ \frac{\mathrm{k}}{\mathrm{k}+1}$ * $ \frac{\mathrm{1}}{\mathrm{k}}$ = $ \frac{\mathrm{1}}{\mathrm{k}+1}$ ，则第 i 个元素不被替换的概率为 1 - $ \frac{\mathrm{1}}{\mathrm{k}+1}$ = $ \frac{\mathrm{k}}{\mathrm{k}+1}$ ，同理，第 k + 2 个元素被选中进行替换的概率为 $ \frac{\mathrm{k}}{\mathrm{k}+2}$ ，第 i 个元素被选中的概率为 $ \frac{\mathrm{1}}{\mathrm{k}}$ ，于是第 i 个元素被第 k + 2 个元素替换的概率为 $ \frac{\mathrm{k}}{\mathrm{k}+2}$ * $ \frac{\mathrm{1}}{\mathrm{k}}$ = $ \frac{\mathrm{1}}{\mathrm{k}+2}$ ，第 i 个元素不被替换的概率为 1 - $ \frac{\mathrm{1}}{\mathrm{k}+2}$ = $ \frac{\mathrm{k}+1}{\mathrm{k}+2} $。以此类推，运行到第 n 步时，被保留的概率 = 被选中的概率 * 不被替换的概率，即：

\[1\times \frac{\mathrm{k}}{\mathrm{k}+1}\times \frac{\mathrm{k}+1}{\mathrm{k}+2}\times \frac{\mathrm{k}+2}{\mathrm{k}+3}\times ...\times \frac{\mathrm{n}-1}{\mathrm{n}}=\frac{\mathrm{k}}{\mathrm{n}} \]

对于第 j 个元素（j > k），该元素被选中的概率为 $ \frac{\mathrm{k}}{\mathrm{j}}$ ，第 j + 1 个元素被选中进行替换的概率为 $ \frac{\mathrm{k}}{\mathrm{j}+1}$，第 j 个元素被选中的概率为 $ \frac{\mathrm{1}}{\mathrm{k}}$，第 j 个元素被替换的概率为 $ \frac{\mathrm{k}}{\mathrm{j}+1}$ * $ \frac{\mathrm{1}}{\mathrm{k}}$ = $ \frac{\mathrm{1}}{\mathrm{j}+1}$，则第 j 个元素不被替换的概率为 1 - $ \frac{\mathrm{1}}{\mathrm{j}+1}$ = $ \frac{\mathrm{j}}{\mathrm{j}+1}$。则运行到第 n 步时，被保留的概率 = 被选中的概率 * 不被替换的概率，即：

\[\frac{\mathrm{k}}{\mathrm{j}}\times \frac{\mathrm{j}}{\mathrm{j}+1}\times \frac{\mathrm{j}+1}{\mathrm{j}+2}\times \frac{\mathrm{j}+2}{\mathrm{j}+3}\times ...\times \frac{\mathrm{n}-1}{\mathrm{n}}=\frac{\mathrm{k}}{\mathrm{n}} \]

所以对于其中每个元素，被保留的概率都为 $ \frac{\mathrm{k}}{\mathrm{n}}$

代码实现

public class ReservoirSampling {
    private int[] pool;  // 蓄水池，包含所有数据
    private int size;    // 蓄水池规格
    private Random random;

    public ReservoirSampling(int size) {
        this.size = size;
        random = new Random();
        // 初始化数据
        pool = new int[size];
        for (int i = 0; i < size; i++) {
            pool[i] = i;
        }
    }

    public int[] sampling(int K) {
        int[] result = new int[K];
        for (int i = 0; i < K; i++) { // 前 K 个元素直接放入数组中
            result[i] = pool[i];
        }

        for (int i = K; i < size; i++) { // K + 1 个元素开始进行概率采样
            int r = random.nextInt(i + 1); // 索引下标为 i 个数据时第 i + 1 个数据，r = [0,i]
            if (r < K) {                         //  选中概率为 k/i+1
                result[r] = pool[i];
            }
        }

        return result;
    }
}

测试

    public static void main(String[] args) {
        ReservoirSampling test = new ReservoirSampling(1000);
        int[] sampling = test.sampling(5);
        for (int i : sampling) {
            System.out.print(i + " ");
        }
    }
// 输出 205 907 986 696 443，每次运行结果不同

题目

LeetCode 382. 链表随机节点

给你一个单链表，随机选择链表的一个节点，并返回相应的节点值。每个节点 被选中的概率一样 。

实现 Solution 类：

Solution(ListNode head) 使用整数数组初始化对象。
int getRandom() 从链表中随机选择一个节点并返回该节点的值。链表中所有节点被选中的概率相等。
 

示例：


输入
["Solution", "getRandom", "getRandom", "getRandom", "getRandom", "getRandom"]
[[[1, 2, 3]], [], [], [], [], []]
输出
[null, 1, 3, 2, 2, 3]

解释
Solution solution = new Solution([1, 2, 3]);
solution.getRandom(); // 返回 1
solution.getRandom(); // 返回 3
solution.getRandom(); // 返回 2
solution.getRandom(); // 返回 2
solution.getRandom(); // 返回 3
// getRandom() 方法应随机返回 1、2、3中的一个，每个元素被返回的概率相等。
 

提示：

链表中的节点数在范围 [1, 104] 内
-104 <= Node.val <= 104
至多调用 getRandom 方法 104 次
 

进阶：

如果链表非常大且长度未知，该怎么处理？
你能否在不使用额外空间的情况下解决此问题？

解：典型的蓄水池算法，当 k 取 1 时的特殊情况，每次只取出一个元素。

class Solution {
    ListNode head;
    Random random = new Random();
    public Solution(ListNode _head) {
        this.head = _head;
    }
    
    // 另第 idx 个结点被选中的概率为 1/idx ，则该结点不被后面结点覆盖的概率为 1/idx * 
    // (1 - 1/(idx+1)) * (1 - 1/(idx+2)) * ...* (1 - 1/n) = 1/n
    // 白话：对第 idx 个结点计算概率，random = [0, idx)， 则random = 0 的概率为 1/idx
    // 只要第 idx 个结点的 random 为 0 则选中覆盖原答案，直到选到最后一个结点。
    public int getRandom() {
        int idx = 1;
        ListNode node = head;
        int ans = node.val;
        while(node != null) {
            if(random.nextInt(idx) == 0) ans = node.val;
            node = node.next;
            idx++;
        }
        return ans;
    }
}

LeetCode 398. 随机数索引

给定一个可能含有重复元素的整数数组，要求随机输出给定的数字的索引。 您可以假设给定的数字一定存在于数组中。

注意：
数组大小可能非常大。 使用太多额外空间的解决方案将不会通过测试。

示例:

int[] nums = new int[] {1,2,3,3,3};
Solution solution = new Solution(nums);

// pick(3) 应该返回索引 2,3 或者 4。每个索引的返回概率应该相等。
solution.pick(3);

// pick(1) 应该返回 0。因为只有nums[0]等于1。
solution.pick(1);

解：只需要考虑给定数字即可，对遍历到的给定数字进行编号（1，2，...），再按照蓄水池算法随机取出一个即可

class Solution {
    private int[] nums;

    public Solution(int[] nums) {
        this.nums = nums;
    }
    
    public int pick(int target) {
        int ans = 0;
        int idx = 0;
        Random random = new Random();
        for(int i = 0; i < nums.length; i++) {
            if(nums[i] == target) {
                idx++;
                if(random.nextInt(idx) == 0) ans = i;
            }
        }
        return ans;
    }   
}

参考资料

蓄水池抽样算法（Reservoir Sampling）
蓄水池采样算法
 挺有意思的一个视频

码农公寓