leetcode 528/ LCR 071 按权重随机选择
题目描述
给定一个正整数数组 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：
输入：
inputs = ["Solution","pickIndex"]
inputs = [[[1]],[]]
输出：
[null,0]
解释：
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。
示例 2：
输入：
inputs = ["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
inputs = [[[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 <= 100001 <= w[i] <= 10^5pickIndex将被调用不超过10000次
解答
第一直觉解答（超时）
根据题目所述：需要随机取下标，但是要保证下标的比例，会很容易想到构建一个按照比例存放的数组，然后随机从中取出下标
例如：
当w = [3, 1, 2]时，我构建一个下标数组：indexArr = [0, 0, 0, 1, 2, 2]，再随机从indexArr 中随机取下标，就能保证0、1、2的比例为w数组所确定的3: 1: 2
代码：
class Solution {
#indexArr: number[] = [];
constructor(w: number[]) {
for (let i = 0; i < w.length; i++) {
this.#indexArr = this.#indexArr.concat(new Array(w[i]).fill(i));
}
}
pickIndex(): number {
return this.#indexArr[Math.round(Math.random() * (this.#indexArr.length - 1))];
}
}
但是此方法在数据量特别大的时候会出现超时，导致无法通过leetcode测试用例：
超时的原因在于concat与fill方法分别为O(n+m)与O(m)的时间复杂度，再在外面套一层遍历w的for循环导致时间复杂度上升到O(n^2)，而本题是有O(n)时间复杂度的解法的，所以这里无法通过测试。
前缀和+二分法
这种方法本质上也是想办法按照比例“拉长”下标数组（例如上述例子，将[3, 1, 2]拉长为[0, 0, 0, 1, 2, 2]），此方法是根据前缀和的特点将[3, 1, 2]“视作”为[1, 2, 3, 4, 5, 6]，再从1至6中取随机数，当随机数为1或2或3时，返回下标0；当随机数为4时，返回下标1；当随机数为5或6时，返回下标2。这样也能保证取出的随机数比例是3: 1: 2。
方法原理：
当w = [3, 1, 2]时，构建前缀和数组[3, 4, 6]，如果此时我们从1开始，在脑海或者草稿纸上使用虚拟的数字填满1至6，可知比例确实是3: 1: 2
代码逻辑
根据上面所述，我们的函数逻辑只需要构建前缀和数组，然后随机取1到前缀和的最后一项的值，然后在前缀和数组中找到刚好大于或等于该随机值的数组元素，然后返回此数组元素对应的下标即可。
class Solution {
#suffixSumArr: number[];
constructor(w: number[]) {
let sum = 0;
// 构建前缀和数组
this.#suffixSumArr = w.map((num) => (sum += num));
}
pickIndex(): number {
const randomNum = Math.random() * this.#suffixSumArr.at(-1)! - 1 + 1;
// 寻找刚好大于randomNum的前缀和数组元素的下标
let i = 0;
while (true) {
if (this.#suffixSumArr[i] >= randomNum) return i;
i++
}
}
}
此时方法还有优化空间，可以采用二分法将“寻找刚好大于随机数的前缀和数组元素的下标”这一步的时间复杂度由O(n)进一步优化到O(logn)
class Solution {
#suffixSumArr: number[];
constructor(w: number[]) {
let sum = 0;
// 构建前缀和数组
this.#suffixSumArr = w.map((num) => (sum += num));
}
pickIndex(): number {
const randomNum = Math.random() * this.#suffixSumArr.at(-1)! - 1 + 1;
let l = 0;
let r = this.#suffixSumArr.length;
// 二分法查找刚好大于或等于randomNum大的值的下标
while (true) {
const middle = (l + r) >> 1;
const middleNum = this.#suffixSumArr[middle];
const middleLeftNum = this.#suffixSumArr[middle - 1] ?? -Infinity;
if (middleNum < randomNum) {
// middle下标对应值过小
l = middle + 1;
continue;
} else if (middleLeftNum >= randomNum) {
// middle下标对应值过大
r = middle - 1;
continue;
}
return middle;
}
}
}