LeetCode 528. 按权重随机选择

约 5 分钟
文章信息 13 个分类 5 个标签
分类
LeetCode CodeTop 中等 未解决 20-50 腾讯 字节 阿里 美团 百度 快手 滴滴 虾皮
标签
数组 数学 二分查找 前缀和 随机化

题目描述

✅ 528. 按权重随机选择

image-20260928204453684

image-20260928204453685

题意分析

给定一组固定的正整数权重 w，每次调用 pickIndex() 随机返回一个下标 i，要求它被选中的概率恰好为 w[i] / sum(w)。

权重表示抽中机会所占的比例，不是返回值；题目要求返回下标。各次调用都按同一分布抽样，不需要扣减权重，也不要求少量调用后的实际次数严格符合比例。

解法：前缀和 + 二分定位随机数

核心思路

[!blue]

如果把下标 i 重复放入数组 w[i] 次，再均匀抽取一个位置，就能得到所需概率，但这样会占用与总权重相关的空间。前缀和可以只记录每段重复区域的边界，不必真的展开这些位置。

令 prefix[i] 为从 w[0] 到 w[i] 的累加和，总和为 total。下标 0 对应 [0, prefix[0])，其余下标 i 对应 [prefix[i - 1], prefix[i])。这些区间互不重叠，完整覆盖 [0, total)，每个区间恰好含有 w[i] 个整数。

每次在 [0, total) 中等概率抽取整数 target，再返回它所在区间的下标。所有整数位置都等概率，落入第 i 段的概率就等于该段整数个数除以总个数，即 w[i] / total。

因为所有权重为正，前缀和严格递增，区间定位就是查找第一个严格大于 target 的前缀和。若 prefix[mid] > target，答案可能就是 mid 或在它左边，保留 mid；否则当前段已经结束，答案必在右边。最终左右边界相遇的位置就是目标下标。

随机值一定小于最后一个前缀和 total，因此答案必然存在。构造时一次完成前缀和，之后每次查询只需一次随机抽样和一次二分。题目约束下总权重不超过 10^9，代码使用的整数类型可以保存。

解题步骤

  1. 构造时从左到右累加权重，保存 prefix 和总权重 total。
  2. 查询时均匀生成满足 0 <= target < total 的整数。
  3. 在前缀和数组中二分：若 prefix[mid] > target，令 right = mid；否则令 left = mid + 1。
  4. 当 left == right 时返回该下标，它对应唯一包含 target 的区间。

代码实现

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
}

复杂度分析

设权重数组长度为 $n$。

  • 时间复杂度：构造为 $O(n)$，单次查询为 $O(\log n)$，仅一个权重时查询为常数时间。
  • 空间复杂度：$O(n)$，保存前缀和；单次查询只使用常数个变量。

关键点总结

[!green]

  • 前缀和隐式表示长度与权重成正比的整数区间，均匀抽位置即按权重抽下标。
  • [0, total) 的随机范围必须与“第一个严格大于目标”的二分条件配套。
  • 正权重保证前缀和严格递增，最后一个前缀和保证二分答案存在。

易错点总结

[!yellow]

  • 使用 [0, total) 抽样却查询 >= target，会把落在段边界上的位置错误归给前一段。
  • 命中 prefix[mid] > target 时不能排除 mid，它本身可能就是答案，应使用 right = mid。
  • 另一分支应使用 left = mid + 1，写成 left = mid 可能无法缩小区间。
  • 不要按权重真的展开数组，空间会随 total 增长，失去前缀和表示的优势。
  • 返回的是区间下标，不是前缀和值；也不能用固定轮流返回的方式代替随机抽样。

相似题目

题目 难度 关联与区别
497. 非重叠矩形中的随机点 中等 都把均匀随机整数放入累计权重区间并二分定位；该题每个矩形的格点数就是权重，选中矩形后再定位其中的格点。
转载与许可
作者
链接 https://hgnulb.github.io/blog/2026/17386994
许可 本博客所有文章除特别声明外，均采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处！