题目描述
这是 LeetCode 上的 528. 按权重随机选择 ，难度为 中等。
Tag : 「前缀和」、「二分」、「模拟」
给定一个正整数数组 w ，其中 w[i] 代表下标 i 的权重（下标从 0 开始），请写一个函数 pickIndex ，它可以随机地获取下标 i，选取下标 i 的概率与 w[i] 成正比。
例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25 （即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。
也就是说，选取下标 i 的概率为 w[i] / sum(w) 。
示例 1：
输入： ["Solution","pickIndex"] [[[1]],[]] 输出： [null,0] 解释： Solution solution = new Solution([1]); solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。 复制代码
示例 2：
输入： ["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"] [[[1,3]],[],[],[],[],[]] 输出： [null,1,1,1,1,0] 解释： Solution solution = new Solution([1, 3]); solution.pickIndex(); // 返回 1，返回下标 1，返回该下标概率为 3/4 。 solution.pickIndex(); // 返回 1 solution.pickIndex(); // 返回 1 solution.pickIndex(); // 返回 1 solution.pickIndex(); // 返回 0，返回下标 0，返回该下标概率为 1/4 。 由于这是一个随机问题，允许多个答案，因此下列输出都可以被认为是正确的: [null,1,1,1,1,0] [null,1,1,1,1,1] [null,1,1,1,0,0] [null,1,1,1,0,1] [null,1,0,1,0,0] ...... 诸若此类。 复制代码
提示：
- 1 <= w.length <= 10000
- 1 <= w[i] <= 10^5105
- pickIndex 将被调用不超过 10000 次
前缀和 + 二分
根据题意，权重值 w[i]w[i] 可以作为 pickIndex 调用总次数为 \sum_{i = 0}^{w.length - 1} w[i]∑i=0w.length−1w[i] 时，下标 ii 的返回次数。
随机数的产生可以直接使用语言自带的 API，剩下的我们需要构造一个分布符合权重的序列。
由于 1 <= w[i] <= 10^51<=w[i]<=105，且 ww 长度为 10^4104，因此直接使用构造一个有 w[i]w[i] 个的 ii 的数字会 MLE。
我们可以使用「前缀和」数组来作为权重分布序列，权重序列的基本单位为 11。
一个长度为 nn 的构造好的「前缀和」数组可以看是一个基本单位为 11 的 [1, sum[n - 1]][1,sum[n−1]] 数轴。
使用随机函数参数产生 [1, sum[n - 1]][1,sum[n−1]] 范围内的随机数，通过「二分」前缀和数组即可找到分布位置对应的原始下标值。
代码：
class Solution { int[] sum; public Solution(int[] w) { int n = w.length; sum = new int[n + 1]; for (int i = 1; i <= n; i++) sum[i] = sum[i - 1] + w[i - 1]; } public int pickIndex() { int n = sum.length; int t = (int) (Math.random() * sum[n - 1]) + 1; int l = 1, r = n - 1; while (l < r) { int mid = l + r >> 1; if (sum[mid] >= t) r = mid; else l = mid + 1; } return r - 1; } } 复制代码
- 时间复杂度：
Solution类的构造方法整体复杂度为 O(n)O(n)；pickIndex的复杂度为 O(\log{n})O(logn) - 空间复杂度：O(n)O(n)
模拟（桶轮询）
利用 OJ 不太聪明（对权重分布做近似检查），我们可以构造一个最小轮询序列（权重精度保留到小数点一位），并使用 (i, cnt)(i,cnt) 的形式进行存储，代表下标 ii 在最小轮询序列中出现次数为 cntcnt。
然后使用两个编号 bidbid 和 iidiid 来对桶进行轮询返回（循环重置 & 跳到下一个桶）。
该解法的最大好处是不需要使用 random 函数，同时返回的连续序列满足每一段（长度不短于最小段）都符合近似权重分布。
代码：
class Solution { // 桶编号 / 桶内编号 / 总数 int bid, iid, tot; List<int[]> list = new ArrayList<>(); public Solution(int[] w) { int n = w.length; double sum = 0, min = 1e9; for (int i : w) { sum += i; min = Math.min(min, i); } double minv = min / sum; int k = 1; while (minv * k < 1) k *= 10; for (int i = 0; i < n; i++) { int cnt = (int)(w[i] / sum * k); list.add(new int[]{i, cnt}); tot += cnt; } } public int pickIndex() { if (bid >= list.size()) { bid = 0; iid = 0; } int[] info = list.get(bid); int id = info[0], cnt = info[1]; if (iid >= cnt) { bid++; iid = 0; return pickIndex(); } iid++; return id; } } 复制代码
- 时间复杂度：计算 kk 的操作只会发生一次，可以看作是一个均摊到每个下标的常数计算，
Solution类的构造方法的整体复杂度可看作 O(n)O(n)；pickIndex的复杂度为 O(1)O(1) - 空间复杂度：O(n)O(n)
最后
这是我们「刷穿 LeetCode」系列文章的第 No.528 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。
在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。
为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：github.com/SharingSour… 。
在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。
