按权重随机选择 图解题解

按权重随机返回下标，pickIndex 要做到什么

给正整数数组 w，w[i] 是下标 i 的权重。pickIndex() 每次随机返回一个下标，返回 i 的概率等于 w[i] 除以权重之和。w=[1,3,2] 时概率是 1/6、3/6、2/6，权重 3 的下标 1 最常被抽中。

把权重摊成一个大数组，问题出在哪

直觉是按权重摊开：1 个 0、3 个 1、2 个 2，成 [0,1,1,1,2,2]，均匀挑下标，抽样 O(1)，但长度等于权重之和。权重上到几万，空间 O(总权重) 就把内存吃爆。得找不真摊开又按权重的办法。

不摊开数组，怎么还能按权重抽

看成一条长度 6 的数轴：下标 0 占 (0,1]、1 占 (1,4]、2 占 (4,6]，每段长度正是那下标的权重，权重越大区间越宽。均匀掷一点，落进哪段就选哪个下标，等于按权重选。

边界用前缀和记（累加权重，pre[i] 为前 i+1 个的和）：pre=[1,4,6]，pre[i] 是下标 i 那段右端。合法的点是 1 到 6 这 6 个整数，掷出 target 就问它落进哪段。pre 升序，找落进的段即找第一个 ≥ target 的位置（bisect_left）。右端是闭区间，target 等于某 pre[i] 时该归下标 i。

bisect_left 在 pre 上怎么一步步夹出下标

bisect_left 找下界：lo、hi 圈候选，起初 lo=0、hi=n，答案在 [lo,hi)。取 mid=(lo+hi)//2，pre[mid] < target 就 lo=mid+1，否则 hi=mid（保留 mid、不减一，因为左边或许还有更小的也满足）。lo、hi 相撞即得下标。

手算 w=[1,3,2]：target=4、5 各夹几轮

前缀和 pre[0]=1，pre[1]=1+3=4，pre[2]=4+2=6，即 pre=[1,4,6]，total=6。掷 target=4，lo=0、hi=3。第一轮 mid=(0+3)//2=1，pre[1]=4，4 < 4 不成立，hi=1，[0,1)。第二轮 mid=0，pre[0]=1 < 4，lo=1。lo=hi=1 停，返回 1，target=4 落在下标 1 的 (1,4]。

再掷 target=5，lo=0、hi=3。第一轮 mid=1，pre[1]=4 < 5，lo=2。第二轮 [2,3)，mid=2，pre[2]=6 ≥ 5，hi=2，返回 2，target=5 落在 (4,6]。把 1 到 6 全掷一遍：1 点归下标 0、3 点归下标 1、2 点归下标 2，比例 1:3:2，正好是权重比。

随机数从 0 起、二分找成 >，概率就偏了

坑不在崩，而在悄悄算偏概率。target 要掷 [1,total]，从 0 起会让多出的 0 挤进第一段，把下标 0 的概率抬高一个点；且必须用 bisect_left 找「第一个 ≥」，换成 bisect_right 找「第一个 >」，target 压在某 pre[i] 上时会算到下标 i+1，边界整段错位。两处都编译得过、跑得动，抽样几万次才看出偏了。

Python

import random, bisect
class Solution:
    def __init__(self, w):
        self.pre = []
        s = 0
        for x in w:                       # 构建前缀和
            s += x
            self.pre.append(s)
        self.total = s
    def pickIndex(self):
        target = random.randint(1, self.total)   # [1, total]
        return bisect.bisect_left(self.pre, target)  # 第一个 ≥

C++

class Solution {
    vector<int> pre;
public:
    Solution(vector<int>& w) {
        int s = 0;
        for (int x : w) { s += x; pre.push_back(s); }  // 前缀和
    }
    int pickIndex() {
        int target = rand() % pre.back() + 1;          // [1, total]
        return lower_bound(pre.begin(), pre.end(), target) - pre.begin();
    }
};

Java

class Solution {
    private int[] pre;
    private java.util.Random rnd = new java.util.Random();
    public Solution(int[] w) {
        pre = new int[w.length];
        int s = 0;
        for (int i = 0; i < w.length; i++) {           // 前缀和
            s += w[i];
            pre[i] = s;
        }
    }
    public int pickIndex() {
        int target = rnd.nextInt(pre[pre.length - 1]) + 1;  // [1, total]
        int lo = 0, hi = pre.length;
        while (lo < hi) {                              // 找第一个 ≥
            int mid = lo + (hi - lo) / 2;
            if (pre[mid] >= target) hi = mid;
            else lo = mid + 1;
        }
        return lo;
    }
}