528. 按权重随机选择
题目描述
设计类 Solution，按权重随机抽取下标。
Solution(w)：用正整数数组 w 初始化，w[i] 是下标 i 的权重，下标从 0 开始。
pickIndex()：随机返回 [0, w.length - 1] 中的一个下标，抽到下标 i 的概率必须等于 w[i] / sum(w)。例如 w = [1,3] 时，返回 0 的概率是 1/4，返回 1 的概率是 3/4。
由于结果随机，只要概率分布正确，任何一组输出都可以被接受。
示例 1
输入：操作 = ["Solution","pickIndex"], 参数 = [[[1]],[]] 输出：[null, 0]
解释：只有一个下标，每次都只能返回 0。
示例 2
输入：操作 = ["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"] 参数 = [[[1,3]],[],[],[],[],[]] 输出：[null, 0, 1, 1, 1, 1]
解释：下标 1 的概率为 3/4，下标 0 为 1/4；这只是可能的输出之一。
数据范围
1 <= w.length <= 10^41 <= w[i] <= 10^5pickIndex最多调用10^4次。
解题思路：前缀和 + 二分查找
把权重拼成一条长度为 sum(w) 的线段：下标 i 占 (pre[i-1], pre[i]] 这 w[i] 个整数点，pre 是前缀和。均匀随机取 x ∈ [1, sum]，落在哪一段就返回哪个下标。pre 严格递增，“第一个 pre[i] >= x 的 i”就是 bisect_left(pre, x)。
为什么正确
每个下标 i 恰好对应 w[i] 个 x 值，所以被选中概率是 w[i] / sum。
复杂度分析
构造 O(n)，每次 pickIndex O(log n)，空间 O(n)。
易错点
随机数范围和查找方式要配套：x 取 [1, sum] 配 bisect_left；若取 [0, sum) 就要配 bisect_right。配错会让某个下标多一个或少一个点，甚至越界。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。
class Solution:
def __init__(self, w: List[int] = ()):
self.pre = []
s = 0
for x in w:
s += x
self.pre.append(s)
def pickIndex(self) -> int:
x = random.randint(1, self.pre[-1])
return bisect_left(self.pre, x)class Solution {
private int[] pre;
private java.util.Random rng = 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 x = rng.nextInt(pre[pre.length - 1]) + 1;
int lo = 0, hi = pre.length; while (lo < hi) { int m = (lo + hi) >>> 1; if (pre[m] < x) lo = m + 1; else hi = m; } return lo;
}
}class Solution {
vector<int> pre;
mt19937 rng{random_device{}()};
public:
Solution(vector<int>& w) {
int s = 0;
for (int x : w) {
s += x;
pre.push_back(s);
}
}
int pickIndex() {
int x = uniform_int_distribution<int>(1, pre.back())(rng);
return lower_bound(pre.begin(), pre.end(), x) - pre.begin();
}
};type Solution struct {
pre []int
}
func Constructor(w []int) Solution {
pre := []int{}
s := 0
for _, x := range w {
s += x
pre = append(pre, s)
}
return Solution{pre}
}
func (this *Solution) PickIndex() int {
x := rand.Intn(this.pre[len(this.pre)-1]) + 1
return sort.SearchInts(this.pre, x)
}var Solution = function(w) {
this.pre = [];
let s = 0;
for (const x of w) {
s += x;
this.pre.push(s);
}
};
Solution.prototype.pickIndex = function() {
const x = 1 + Math.floor(Math.random() * this.pre[this.pre.length - 1]);
let lo = 0, hi = this.pre.length; while (lo < hi) { const m = (lo + hi) >> 1; if (this.pre[m] < x) lo = m + 1; else hi = m; } return lo;
};class Solution {
pre: number[];
constructor(w: number[]) {
this.pre = [];
let s = 0;
for (const x of w) {
s += x;
this.pre.push(s);
}
}
pickIndex(): number {
const x = 1 + Math.floor(Math.random() * this.pre[this.pre.length - 1]);
let lo = 0, hi = this.pre.length; while (lo < hi) { const m = (lo + hi) >> 1; if (this.pre[m] < x) lo = m + 1; else hi = m; } return lo;
}
}其他解法
前缀和加二分把权重铺在数轴上，每次采样做一次二分；别名法预处理成 n 个“最多两种结果”的桶，把单次采样降到 O(1)。
解法 2：别名法（Vose Alias Method） 时间 初始化 O(n)，每次选取 O(1)，空间 O(n)
把每个下标的概率乘以 n，平均值变为 1。准备 n 个容量为 1 的桶：每次取一个不足 1 的下标 s 放进自己的桶，用一个不少于 1 的下标 l 把桶补满，记录 prob[s] 和别名 alias[s] = l，l 的剩余量相应减少，再按是否小于 1 放回对应列表。选取时先等概率选一个桶 i，再以 prob[i] 的概率返回 i，否则返回 alias[i]。
每次选取只需两次随机数，时间 O(1)，适合大量采样。构建过程较难理解，并涉及浮点数，存在微小的精度误差。
class Solution:
def __init__(self, w: List[int] = ()):
# Vose 别名法：把每个下标的概率放大 n 倍后，切成 n 个容量为 1 的桶，
# 每个桶最多装两个下标：自己（概率 prob[i]）和别名 alias[i]
n = len(w)
total = sum(w)
self.n = n
self.prob = [0.0] * n
self.alias = list(range(n))
scaled = [x * n / total for x in w]
small = [i for i, p in enumerate(scaled) if p < 1]
large = [i for i, p in enumerate(scaled) if p >= 1]
while small and large:
s, l = small.pop(), large.pop()
self.prob[s] = scaled[s]
self.alias[s] = l # 桶 s 的剩余部分由 l 填满
scaled[l] -= 1 - scaled[s]
(small if scaled[l] < 1 else large).append(l)
for i in small + large: # 剩下的桶恰好装满（浮点误差处理）
self.prob[i] = 1.0
def pickIndex(self) -> int:
# 等概率选一个桶，再以 prob 决定取桶主人还是别名
i = random.randrange(self.n)
return i if random.random() < self.prob[i] else self.alias[i]class Solution {
private int n;
private double[] prob;
private int[] alias;
private java.util.Random rng = new java.util.Random();
public Solution(int[] w) {
n = w.length;
double total = 0;
for (int x : w) total += x;
prob = new double[n];
alias = new int[n];
double[] scaled = new double[n];
Deque<Integer> small = new ArrayDeque<>(), large = new ArrayDeque<>();
for (int i = 0; i < n; i++) {
prob[i] = 1.0;
alias[i] = i;
scaled[i] = w[i] * n / total;
if (scaled[i] < 1) small.push(i);
else large.push(i);
}
while (!small.isEmpty() && !large.isEmpty()) {
int s = small.pop(), l = large.pop();
prob[s] = scaled[s];
alias[s] = l;
scaled[l] -= 1 - scaled[s];
if (scaled[l] < 1) small.push(l);
else large.push(l);
}
}
public int pickIndex() {
int i = rng.nextInt(n);
return rng.nextDouble() < prob[i] ? i : alias[i];
}
}class Solution {
int n;
vector<double> prob;
vector<int> alias;
mt19937 rng{random_device{}()};
public:
Solution(vector<int>& w) {
n = w.size();
double total = 0;
for (int x : w) total += x;
prob.assign(n, 1.0);
alias.resize(n);
vector<double> scaled(n);
vector<int> small, large;
for (int i = 0; i < n; i++) {
alias[i] = i;
scaled[i] = w[i] * n / total;
(scaled[i] < 1 ? small : large).push_back(i);
}
while (!small.empty() && !large.empty()) {
int s = small.back(), l = large.back();
small.pop_back();
large.pop_back();
prob[s] = scaled[s];
alias[s] = l;
scaled[l] -= 1 - scaled[s];
(scaled[l] < 1 ? small : large).push_back(l);
}
}
int pickIndex() {
int i = uniform_int_distribution<int>(0, n - 1)(rng);
double r = uniform_real_distribution<double>(0, 1)(rng);
return r < prob[i] ? i : alias[i];
}
};type Solution struct {
prob []float64
alias []int
}
func Constructor(w []int) Solution {
n := len(w)
total := 0
for _, x := range w {
total += x
}
prob := make([]float64, n)
alias := make([]int, n)
scaled := make([]float64, n)
small, large := []int{}, []int{}
for i, x := range w {
prob[i] = 1.0
alias[i] = i
scaled[i] = float64(x) * float64(n) / float64(total)
if scaled[i] < 1 {
small = append(small, i)
} else {
large = append(large, i)
}
}
for len(small) > 0 && len(large) > 0 {
s, l := small[len(small)-1], large[len(large)-1]
small, large = small[:len(small)-1], large[:len(large)-1]
prob[s] = scaled[s]
alias[s] = l
scaled[l] -= 1 - scaled[s]
if scaled[l] < 1 {
small = append(small, l)
} else {
large = append(large, l)
}
}
return Solution{prob, alias}
}
func (this *Solution) PickIndex() int {
i := rand.Intn(len(this.prob))
if rand.Float64() < this.prob[i] {
return i
}
return this.alias[i]
}var Solution = function(w) {
const n = w.length;
const total = w.reduce((a, b) => a + b, 0);
this.prob = new Array(n).fill(1);
this.alias = w.map((_, i) => i);
const scaled = w.map(x => x * n / total);
const small = [], large = [];
scaled.forEach((p, i) => (p < 1 ? small : large).push(i));
while (small.length && large.length) {
const s = small.pop(), l = large.pop();
this.prob[s] = scaled[s];
this.alias[s] = l;
scaled[l] -= 1 - scaled[s];
(scaled[l] < 1 ? small : large).push(l);
}
};
Solution.prototype.pickIndex = function() {
const i = Math.floor(Math.random() * this.prob.length);
return Math.random() < this.prob[i] ? i : this.alias[i];
};class Solution {
prob: number[];
alias: number[];
constructor(w: number[]) {
const n = w.length;
const total = w.reduce((a, b) => a + b, 0);
this.prob = new Array(n).fill(1);
this.alias = w.map((_, i) => i);
const scaled = w.map(x => x * n / total);
const small: number[] = [], large: number[] = [];
scaled.forEach((p, i) => (p < 1 ? small : large).push(i));
while (small.length && large.length) {
const s = small.pop()!, l = large.pop()!;
this.prob[s] = scaled[s];
this.alias[s] = l;
scaled[l] -= 1 - scaled[s];
(scaled[l] < 1 ? small : large).push(l);
}
}
pickIndex(): number {
const i = Math.floor(Math.random() * this.prob.length);
return Math.random() < this.prob[i] ? i : this.alias[i];
}
}