528. 按权重随机选择
528. 按权重随机选择
🟠 🔖 数组
数学
二分查找
前缀和
随机化
🔗 力扣
LeetCode
题目
You are given a 0-indexed array of positive integers w
where w[i]
describes the weight of the ith
index.
You need to implement the function pickIndex()
, which randomly picks an index in the range [0, w.length - 1]
(inclusive) and returns it. The probability of picking an index i
is w[i] / sum(w)
.
- For example, if
w = [1, 3]
, the probability of picking index0
is1 / (1 + 3) = 0.25
(i.e.,25%
), and the probability of picking index1
is3 / (1 + 3) = 0.75
(i.e.,75%
).
Example 1:
Input
["Solution","pickIndex"]
[[[1]],[]]
Output
[null,0]
Explanation
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. The only option is to return 0 since there is only one element in w.
Example 2:
Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output
[null,1,1,1,1,0]
Explanation
Solution solution = new Solution([1, 3]); solution.pickIndex(); // return 1. It is returning the second element (index = 1) that has a probability of 3/4. solution.pickIndex(); // return 1 solution.pickIndex(); // return 1 solution.pickIndex(); // return 1 solution.pickIndex(); // return 0. It is returning the first element (index = 0) that has a probability of 1/4.
Since this is a randomization problem, multiple answers are allowed.
All of the following outputs can be considered correct:
[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] ......
and so on.
Constraints:
1 <= w.length <= 10^4
1 <= w[i] <= 10^5
pickIndex
will be called at most10^4
times.
题目大意
给你一个 下标从 0 开始 的正整数数组 w
,其中 w[i]
代表第 i
个下标的权重。
请你实现一个函数 pickIndex
,它可以 随机地 从范围 [0, w.length - 1]
内(含 0
和 w.length - 1
)选出并返回一个下标。选取下标 i
的 概率 为 w[i] / sum(w)
。
- 例如,对于
w = [1, 3]
,挑选下标0
的概率为1 / (1 + 3) = 0.25
(即,25%),而选取下标1
的概率为3 / (1 + 3) = 0.75
(即,75%
)。
示例 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 <= 10^4
1 <= w[i] <= 10^5
pickIndex
将被调用不超过10^4
次
解题思路
构建前缀和数组:
- 权重数组
w
中的每个值代表选中对应下标的相对概率。 - 我们将权重数组转换为 前缀和数组,
arr[i]
表示权重数组从开始到第i
项的累积和。 - 这样可以将权重范围映射到连续区间上。例如:
- 对于
w = [1, 3, 2]
,前缀和数组为[1, 4, 6]
,其中:- 区间
[1, 1]
对应下标0
。 - 区间
[2, 4]
对应下标1
。 - 区间
[5, 6]
对应下标2
。
- 区间
- 对于
- 权重数组
随机化:
- 生成一个范围在
[1, sum(w)]
的随机数,sum(w)
为权重总和。 - 通过该随机数找到对应的区间,进而确定返回的下标。
- 生成一个范围在
二分查找:
- 在前缀和数组中找到第一个大于等于随机数的位置(下标
i
)。 - 由于前缀和数组是递增的,可以利用二分查找快速定位。
- 在前缀和数组中找到第一个大于等于随机数的位置(下标
复杂度分析
时间复杂度:
- 构造器
Solution(w)
:构建前缀和数组需要遍历权重数组,时间复杂度为O(n)
。 - 方法
pickIndex()
:使用二分查找确定下标,时间复杂度为O(log n)
。
- 构造器
空间复杂度:
O(n)
,存储前缀和数组需要额外的空间。
代码
/**
* @param {number[]} w
*/
var Solution = function (w) {
let sum = 0;
// 构建前缀和数组
this.arr = w.map((num) => {
sum += num;
return sum;
});
this.sum = sum; // 记录权重总和
};
/**
* @return {number}
*/
Solution.prototype.pickIndex = function () {
// 生成 [1, sum] 的随机数
const random = Math.floor(Math.random() * this.sum) + 1;
let left = 0,
right = this.arr.length;
// 二分查找第一个大于等于 random 的位置
while (left < right) {
const mid = Math.floor((left + right) / 2);
if (this.arr[mid] < random) {
left = mid + 1;
} else {
right = mid;
}
}
return left;
};
/**
* Your Solution object will be instantiated and called as such:
* var obj = new Solution(w);
* var param_1 = obj.pickIndex();
*/
相关题目
题号 | 标题 | 题解 | 标签 | 难度 | 力扣 |
---|---|---|---|---|---|
398 | 随机数索引 | 水塘抽样 哈希表 数学 1+ | 🟠 | 🀄️ 🔗 | |
497 | 非重叠矩形中的随机点 | 水塘抽样 数组 数学 4+ | 🟠 | 🀄️ 🔗 | |
710 | 黑名单中的随机数 | 数组 哈希表 数学 3+ | 🔴 | 🀄️ 🔗 |