【二分查找】LeetCode 528. 按权重随机选择

题目链接

528. 按权重随机选择

思路

参考宫水三叶大佬思路

可以回忆一下概率论中所学到的：密度函数与分布函数

在这个题里面呢，数组 w 其实就给出了下标 i 被选中的“概率”（因为总和不是1，所以加了引号），而我们知道离散情况下，分布函数是由密度函数加和求出来的，所以可以使用前缀和 \(sum\) 存储其分布函数。

因为每个数的权重有大有小，所以前缀和中的每个数值间的“空隙”也有大有小，而这个大小就代表了选择这个数字的大小。

之后我们可以使用 Math.random() 产生一个位于 \([0, 1)\) 内的随机数，使用这个随机数乘上 \(sum[n - 1]\) 得到 \(target\)

因为 \(sum\) 数组本身是有序的，所以我们可以使用二分查找来搜索 \(target\)。

代码

class Solution {
    private int[] sum;

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

    public int pickIndex() {
        int n = sum.length;
        // 0 <= random < 1
        // => 0 <= (int)(random * sum[n - 1]) <= sum[n - 1] - 1
        // => 1 <= (int)(random * sum[n - 1]) + 1 <= sum[n - 1]
        int target = (int)(Math.random() * sum[n - 1]) + 1;
        int left = 1;
        int right = n - 1;

        while(left <= right){
            int mid = (right - left) / 2 + left;
            if(sum[mid] == target){
                return mid - 1;
            }else if(sum[mid] > target){
                right = mid - 1;
            }else{
                left = mid + 1;
            }
        }

        return left - 1;
    }
}