LeetCode LCR 071. 按权重随机选择
题目描述
题意分析
给定正整数权重数组
w，每次调用pickIndex()都随机返回一个下标，使下标i被选中的概率为 $w[i]/\sum w$。多次调用使用同一份权重，不会因为某个下标被选过就减少它的权重。可以把每份权重理解为一个可被抽中的整数位置。下标
i占据w[i]个位置，从全部位置中均匀选一个，就能按区间长度得到所需概率。无需真的创建这些位置，只需保存每段的边界。
解法：前缀权重区间与二分采样
核心思路
[!blue]
设权重个数为 $m$，用
presum[i]表示前i项权重之和，初值presum[0] = 0，总权重为T = presum[m]。下标i对应整数区间 $(presum[i],presum[i+1]]$，其中恰好有w[i]个整数。这些区间互不重叠，并完整覆盖
1..T。从中随机选一个落点x后，只需找到第一个满足presum[i+1] >= x的下标i；前一段的右端尚未到达落点，当前段的右端已经到达，落点恰好属于当前段。权重都为正，前缀和严格递增，可以对段号二分。维护包含答案的闭区间，若第
mid段的右端已经不小于x，答案在它或其左侧，令right = mid；否则该段及左侧都无法覆盖落点，令left = mid+1。最后一段的右端是总权重，必定覆盖某个合法落点，因此答案始终存在。概率由分段长度保证：总共有
T个可抽位置，恰有w[i]个位置会被定位到下标i，因此抽中概率与权重成正比。这个结论描述每次抽样的分布，不要求有限次调用的结果严格按权重比例交替出现。前缀和只在构造时计算，后续查询只负责抽点和定位。题目中长度最多 $10^4$、单项权重最多 $10^5$，总和最多 $10^9$，当前整数类型足够保存。
解题步骤
- 创建长度为
w.length+1的前缀和数组，将权重逐项累加。- 每次查询生成
1..T内的随机整数。Java 将[0,1)的随机数缩放取整后加一；Go 用rand.Intn(T)+1。- 在原下标范围
[0,w.length-1]内查找首个右端点不小于落点的段，返回其段号。查询代码里的
n是前缀和数组长度，所以总权重位于presum[n-1]，合法段号的右端是n-2。只有一个权重时，二分直接返回下标零。落点等于某段右端点时仍属于该段，比较条件必须含等号。
代码实现
class Solution {
private int[] presum;
public Solution(int[] w) {
int n = w.length;
presum = new int[n + 1];
for (int i = 0; i < n; ++i) {
presum[i + 1] = presum[i] + w[i];
}
}
public int pickIndex() {
int n = presum.length;
int x = (int) (Math.random() * presum[n - 1]) + 1;
int left = 0;
int right = n - 2;
while (left < right) {
int mid = (left + right) >> 1;
if (presum[mid + 1] >= x) {
right = mid;
} else {
left = mid + 1;
}
}
return left;
}
}
import (
"math/rand"
)
type Solution struct {
presum []int
}
func Constructor(w []int) Solution {
n := len(w)
pre := make([]int, n+1)
for i := 0; i < n; i++ {
pre[i+1] = pre[i] + w[i]
}
return Solution{pre}
}
func (this *Solution) PickIndex() int {
n := len(this.presum)
x := rand.Intn(this.presum[n-1]) + 1
left, right := 0, n-2
for left < right {
mid := (left + right) >> 1
if this.presum[mid+1] >= x {
right = mid
} else {
left = mid + 1
}
}
return left
}
复杂度分析
- 时间复杂度：构造为 $O(m)$，每次查询为 $O(\log(m+1))$，由一次随机取点和一次二分定位组成。
- 空间复杂度：$O(m)$，只保存前缀和，不按总权重展开下标。
关键点总结
[!green]
- 权重决定区间中的整数数量，均匀抽点后按区间定位即可实现按权重抽样。
- 本实现采用左开右闭区间，随机落点必须与之配套，取
[1,T]。- 前缀和位置比原下标多一位，第
i段的右端是presum[i+1]，返回值仍是i。
易错点总结
[!yellow]
- 抽到零或超过总权重，会让落点落在全部合法区间之外。
- 把二分条件写成严格大于，会把恰好位于右端点的落点错误分给下一段。
- 每次查询重新建立前缀和，会重复执行可以一次完成的预处理。
- 每次抽样后扣减权重，会改变后续调用应使用的概率分布。
相似题目
| 题目 | 难度 | 关联与区别 |
|---|---|---|
| 497. 非重叠矩形中的随机点 | 中等 | 该题以矩形格点数作为权重，复用前缀和与二分抽样，再把区间内的偏移量转换为格点坐标。 |