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 index 0 is 1 / (1 + 3) = 0.25 (i.e., 25%), and the probability of picking index 1 is 3 / (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 most 10^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 次

解题思路

1. 构建前缀和数组：

   • 权重数组 w 中的每个值代表选中对应下标的相对概率。
   • 我们将权重数组转换为 前缀和数组，arr[i] 表示权重数组从开始到第 i 项的累积和。
   • 这样可以将权重范围映射到连续区间上。例如：
     • 对于 w = [1, 3, 2]，前缀和数组为 [1, 4, 6]，其中：
       • 区间 [1, 1] 对应下标 0。
       • 区间 [2, 4] 对应下标 1。
       • 区间 [5, 6] 对应下标 2。

2. 随机化：

   • 生成一个范围在 [1, sum(w)] 的随机数，sum(w) 为权重总和。
   • 通过该随机数找到对应的区间，进而确定返回的下标。

3. 二分查找：

   • 在前缀和数组中找到第一个大于等于随机数的位置（下标 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+	🔴	🀄️ 🔗