水塘抽样算法

数据结构与算法

发布日期: 2022-04-25

文章字数: 1.2k

阅读时长: 4 分

水塘抽样算法

这个算法是在做leetcode的2022.4.25每日一题学习到的

给你一个可能含有重复元素的整数数组 nums ，请你随机输出给定的目标数字 target 的索引。你可以假设给定的数字一定存在于数组中。

实现 Solution 类：

Solution(int[] nums) 用数组 nums 初始化对象。

int pick(int target) 从 nums 中选出一个满足 nums[i] == target 的随机索引 i 。如果存在多个有效的索引，则每个索引的返回概率应当相等。

最开始自己的想法就是用一个map保存num中的每个元素对应的索引，如果相同元素个数超过一个的话，就用Random函数随机选择索引。但是这种方法的空间复杂度是O(n)，如果数据量特别大的时候（即大数据环境下，这种方法可能性能不太好，故有了水塘抽样的方法。

算法描述

遍历 nums，当我们第 i 次遇到值为 target 的元素时，随机从区间[0,i) 内选择一个随机的整数，如果其等于0，则将返回值置为该元素的下标，否则返回值不变。

算法正确性验证

设 nums 中有 k 个值为 target 的元素，该算法会保证这 k 个元素的下标成为最终返回值的概率均为 1/k，证明如下：

因此，对于这个nums数组中的所有元素而言，每个元素的下标作为最终结果的概率都是一样的，故其实题目也可改为获取随机一个元素的下标。上述问题的代码可以写为：

class Solution {
    int[] nums;
    Random random;

    public Solution(int[] nums) {
        this.nums = nums;
        random = new Random();
    }

    public int pick(int target) {
        int ans = 0;
        for (int i = 0, cnt = 0; i < nums.length; ++i) {
            if (nums[i] == target) {
                ++cnt; // 第 cnt 次遇到 target
                if (random.nextInt(cnt) == 0) {
                    ans = i;
                }
            }
        }
        return ans;
    }
}

算法一般化

上述题目只是选取一个元素的下标作为最终结果，假设我们要选取k个元素的下标，那这样又要怎么实现呢？

其实思路跟上面差不多，就是用一个数组记录下每个要选取元素的cnt即可，也就是说：

public int[] pick(int[] targets) {
    int[] ans = new int[targets.length];
    int[] cnt = new int[targets.length];
    for (int i = 0; i < nums.length; ++i) {
        for(int j = 0; j < targets.length; j++){
            int target = targets[j];
            if (nums[i] == target) {
                ++cnt[j]; // 第 cnt 次遇到 target
                if (random.nextInt(cnt[j]) == 0) {
                    ans[j] = i;
                }
                break;
            }
        }
    }
    return ans;
}

更一般化的结论

当内存无法加载全部数据时，如何从包含未知大小的数据流中随机选取k个数据，并且要保证每个数据被抽取到的概率相等

根据上面算法的类似思想，首先$k=1$时，我们可以看成是当第i个随机得到的数为0时，才保留第i个数的索引，也就是说，随机选择有i种可能出现的情况，只选择其中出现的一种情况，则以此可推知，当这个数据流的第i个数被保留的概率设置为**$1/i$**时，能够保证所有数据的选取概率都是1/N

这样，推及k>1时，我们可以得到这样的结论：

对于前k个数，我们都保留，对于第i（i > k）个数，我们以$k/i$的概率保留这个数，并以$1/k$的概率与之前保留的k个数种的任意一个替换，则选择每一个数的概率都会是1/N

参考代码：

public void ReservoirSampling(int[] results, int[] nums, int k)
{
    int N = nums.size();
    Random random = new Random();

    for (int i=0; i<k; ++i) {
        results[i] = nums[i];
    }

    for (int i=k; i<N; ++i) {
        int random = random.nextInt(i+1);
        if (random<k) {
            results[random] = nums[i];
        }
    }

    return results;
}