LeetCode 528. 按权重随机选择
目录
题目描述
题意分析
构造函数收到一个正整数数组
w，之后pickIndex会被反复调用，每次要返回一个下标i，且返回i的概率必须严格等于w[i]占总和的比例。判题不是比对某一次的返回值，而是统计大量调用后的频率分布是否符合预期。约束里的信号很明确：
w一旦给定就不再变化，而pickIndex的调用次数远大于数组长度，说明应该在构造阶段把重活干完，把单次查询压到尽可能低；权重都是正数，保证累加得到的序列严格递增，这一点是后面能做查找的前提；权重之和可能很大，所以不能按权重把下标一个个展开成一张大表。边界有三处：数组只有一个元素时必须恒定返回 0；某些权重可能远大于其他权重，几乎所有调用都应落在它身上；随机数的取值范围必须与「区间」的定义严丝合缝，差一就会让首尾两个下标的概率偏离。
解法：前缀和 + 二分定位随机数
核心思路
把每个权重看成一段连续区间的长度。构造前缀和
prefix，其中prefix[i]是w[0]到w[i]的总和，total是所有权重之和。在整数集合
[0, total)中等概率生成target。令prefix[-1]代表0，下标i对应区间[prefix[i - 1], prefix[i])，其中恰好包含w[i]个整数。因此target落入该区间的概率为 $w[i] / total$，正好满足题意。
prefix严格递增，所以目标下标就是第一个满足prefix[i] > target的位置，可用二分查找。二分过程中始终保证答案位于[left, right]；命中条件时保留mid，否则排除mid及其左侧。这比按权重展开下标数组更合适：构造空间只与元素个数有关，不与权重总和有关。
解题步骤
- 构造时计算前缀和数组和总权重
total。- 每次调用在
[0, total)中均匀生成整数target。- 在
prefix中二分查找第一个严格大于target的位置。- 返回该位置作为抽中的下标。
例如
w = [1, 3]，前缀和是[1, 4]。随机数0属于下标0，随机数1、2、3属于下标1；四个随机数等概率，因此两个下标的概率分别为 $1/4$ 和 $3/4$。边界上，
target == prefix[i]应归入下一个区间，这正是二分条件使用“严格大于”的原因。
代码实现
class Solution {
private final int[] prefix;
private final int total;
public Solution(int[] w) {
prefix = new int[w.length];
int sum = 0;
for (int i = 0; i < w.length; i++) {
sum += w[i];
prefix[i] = sum;
}
total = sum;
}
public int pickIndex() {
int target = java.util.concurrent.ThreadLocalRandom
.current()
.nextInt(total);
int left = 0;
int right = prefix.length - 1;
while (left < right) {
int mid = left + (right - left) / 2;
if (prefix[mid] > target) {
right = mid;
} else {
left = mid + 1;
}
}
return left;
}
}
import "math/rand"
type Solution struct {
prefix []int
total int
}
func Constructor(w []int) Solution {
prefix := make([]int, len(w))
sum := 0
for i, weight := range w {
sum += weight
prefix[i] = sum
}
return Solution{prefix: prefix, total: sum}
}
func (this *Solution) PickIndex() int {
target := rand.Intn(this.total)
left, right := 0, len(this.prefix)-1
for left < right {
mid := left + (right-left)/2
if this.prefix[mid] > target {
right = mid
} else {
left = mid + 1
}
}
return left
}
复杂度分析
- 构造时间复杂度：$O(n)$。
- 单次查询时间复杂度：$O(\log n)$。
- 空间复杂度：$O(n)$，用于保存前缀和数组。
关键点总结
- 权重对应区间长度，均匀随机点落入某段的概率就与该段长度成正比。
- 随机范围、区间开闭和二分条件必须成套：本实现使用
[0, total)、左闭右开区间和“第一个前缀和大于目标”。- 权重为正数，因此前缀和严格递增，可以二分。
- 权重固定、查询频繁时，适合把 $O(n)$ 的工作放到构造阶段。
- 若面试追问动态修改权重，可使用树状数组，将更新和查询都控制在 $O(\log n)$。
易错点总结
- 随机数范围与二分条件不一致：使用
[0, total)时应查找第一个prefix[i] > target。- 把条件写成
prefix[mid] >= target：w = [1, 3]、target = 1时会错误返回下标0，而它应属于下标1。- 命中后写
right = mid - 1：mid本身可能就是答案，不能排除。- 未命中后写
left = mid：两指针相邻时无法继续收缩，造成死循环。- 按权重展开数组：空间复杂度变成 $O(total)$，大权重下不可接受。
- 返回前缀和值而不是数组下标：题目要求返回的是
i。
相似题目
| 题目 | 难度 | 考察点 |
|---|---|---|
| 382. 链表随机节点 | 中等 | 长度未知时的蓄水池抽样 |
| 398. 随机数索引 | 中等 | 重复值中等概率挑一个下标 |
| 470. 用 Rand7() 实现 Rand10() | 中等 | 由已有均匀源构造新均匀源 |
| 497. 非重叠矩形中的随机点 | 中等 | 先按面积选矩形再选点 |
| LCR 071. 按权重随机选择 | 中等 | 同一模型的换皮版本 |