跳至主要內容

528. 按权重随机选择


528. 按权重随机选择

🟠   🔖  数组 数学 二分查找 前缀和 随机化  🔗 力扣open in new window LeetCodeopen in new window

题目

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] 内(含 0w.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+🟠🀄️open in new window 🔗open in new window
497非重叠矩形中的随机点水塘抽样 数组 数学 4+🟠🀄️open in new window 🔗open in new window
710黑名单中的随机数数组 哈希表 数学 3+🔴🀄️open in new window 🔗open in new window