力扣 528. 按权重随机选择

题目来源：https://leetcode-cn.com/problems/random-pick-with-weight/

大致题意：
给定一个数组，数组索引 i 位置的值 w[i] 表示该索引在数组整体中的权重。设计一个方法，随机返回一个索引，被选中的概率与该位置权重相关。

思路

前缀和 + 二分

若所有权重的总值为 sum，权重个数为 n，可以从 1 到 sum 划分成 n 个区间。使用 pre 表示索引 i 之前的所有权重和（包括），那么 [ pre - w[i] + 1, pre ] 即是索引 i 对应的区间。
于是可以生成一个在区间 [1, sum] 之间的随机数，将其对应的区间的索引返回。

若使用 pre[i] 表示包含索引 i 的前缀和，那么可知前缀和数组 pre 是递增的。于是可以使用二分查找的方法，找出最大的且不大于随机数 x 的 pre[i]，i 即是所求索引。

代码：

public class RandomPick {
    int[] pre;
    int sum;
    public RandomPick(int[] w) {
        // 初始化前缀和数组
        pre = new int[w.length];
        pre[0] = w[0];
        for (int i = 1; i < w.length; i++) {
            pre[i] = pre[i-1] + w[i];
        }
        sum = Arrays.stream(w).sum();
    }

    public int pickIndex() {
        // 生成 [1, sum] 之间的随机数 x
        int x = (int)(Math.random() * sum) + 1;
        // 二分查找
        int idx = binarySearch(x);
        return idx;
    }

    public int binarySearch(int target) {
        int left = 0;
        int right = pre.length - 1;
        while (left < right) {
            int mid = (left + right) / 2;
            // target 大于当前位置前缀和，向右半段查找
            if (pre[mid] < target) {
                left = mid + 1;
            }
            // target 小于等于当前位置前缀和，向左半段查找
            else {
                right = mid;
            }
        }
        return left;
    }
}

收获

快速的生成随机数的方法

生成 [1, sum] 之间随机数的方法

int x = (int)(Math.random() * sum) + 1;