LeetCode 528. 按权重随机选择

题目链接：

力扣https://leetcode-cn.com/problems/random-pick-with-weight/

【分析】如果直接用水塘取样的话会超时，因为查询次数为10^4，每次查询时需要枚举10^4长度的数组，而Random.nextInt()中的参数可以达到10^9。

class Solution {

    Random random = new Random();
    int[] w;
    int[] p;
    int n;

    public Solution(int[] w) {
        this.w = w;
        this.n = w.length;
        p = new int[n];
        p[0] = w[0];
        for(var i = 1; i < n; i++) p[i] = p[i - 1] + w[i];
    }
    
    public int pickIndex() {
        int ans = 0;
        for(var i = 0; i < n; i++){
            if(random.nextInt(p[i]) < w[i]) ans = i;
        }
        return ans;
    }
}

【方法二 前缀和+枚举】把概率值求前缀和，然后生成一个概率综合范围内的随机数r，查找到第一个>r的前坠值，也就把概率用区间长度来表示了，看随机数落在哪个区间内。

class Solution {

    int[] w;
    int n;

    public Solution(int[] w) {
        this.w = w;
        n = w.length;
        for(var i = 1; i < n; ++i){
            w[i] += w[i - 1];
        }
    }
    
    public int pickIndex() {
        double random = Math.random() * w[n - 1];
        for(var i = 0; i < n; ++i){
            if(w[i] > random) return i;
        }
        return n - 1;
    }
}

【优化 前缀和+二分查找】我们发现前缀和事一个排好序的数组，我们要查找大于目标r的最小值，所以这个查询过程可以用二分来实现。

class Solution {

    int[] w;
    int n;
    Random random = new Random();

    public Solution(int[] w) {
        this.w = w;
        n = w.length;
        for(var i = 1; i < n; ++i){
            w[i] += w[i - 1];
        }
    }

    public int BinarySearch(int left, int right, double target){
        int mid;
        while(left <= right){
            mid = (left + right) / 2;
            if(w[mid] <= target) left = mid + 1;
            else right = mid - 1;
        }
        return left;
    }
    
    public int pickIndex() {
        double r = random.nextInt(w[n - 1]);
        return BinarySearch(0, n - 1, r);
    }
}