LeetCode-Go 题解：528. Random Pick with Weight —— 前缀和 + 二分查找实现权重随机采样

LeetCode-Go 题解：528. Random Pick with Weight —— 前缀和 + 二分查找实现权重随机采样

【免费下载链接】LeetCode-Go✅ Solutions to LeetCode by Go, 100% test coverage, runtime beats 100% | LeetCode 题解项目地址: https://gitcode.com/GitHub_Trending/le/LeetCode-Go

导读

本文围绕 LeetCode-Go 仓库中 528. Random Pick with Weight 一题的官方题解文档展开，深入讲解"按权重随机取下标"这一经典算法问题：先构建权重前缀和数组，再在[1, total]区间内取随机数，通过二分查找定位目标下标，从而让每个下标被选中的概率与其权重严格成正比。读完本文，你将掌握前缀和数组 + 二分查找这一套 O(log n) 的权重随机采样方案，并能够结合仓库中的 Go 源码 与 测试用例 完整理解其实现与正确性验证方法。

题目理解

给定一个正整数数组w，其中w[i]表示下标i的权重。需要实现一个pickIndex()函数，它随机返回一个下标i，并且返回下标i的概率与w[i]成正比。例如权重数组w = [1, 3]表示下标0被选中的概率是 1/4，下标1被选中的概率是 3/4。

题目约束如下：

1. 1 <= w.length <= 10000
2. 1 <= w[i] <= 10^5
3. pickIndex将被调用不超过10000次

即权重数组长度最多 10000，单个权重最大 10^5，总权重可达 10^9，因此中间计算需要使用 64 位整数（Go 的int在 64 位平台上足以覆盖该范围）。

输入输出示例

示例 1：

Input: ["Solution","pickIndex"] [[[1]],[]] Output: [null,0]

只有一个下标，权重为 1，无论调用多少次pickIndex都返回 0。

示例 2：

Input: ["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"] [[[1,3]],[],[],[],[],[]] Output: [null,0,1,1,1,0]

权重数组为[1, 3]，下标1的权重是下标0的三倍，因此多次调用后返回1的次数约为返回0的三倍。示例输出中的0,1,1,1,0正是这种比例关系的体现（实际输出具有随机性，不要求与示例完全一致）。

输入语法说明

输入由两个列表组成：第一个列表是依次调用的成员函数名（Solution构造函数与pickIndex），第二个列表是对应函数调用的参数。Solution的构造函数接收一个参数：权重数组w；pickIndex没有参数。所有参数都被包装在一个列表中，即使某个函数没有参数，也会传入一个空列表[]（如上例中的[],[],[],[],[]）。

解题思路：前缀和 + 二分查找

这是典型的"加权随机采样"（weighted random sampling）问题，其核心思想是把权重分布转化到一维数轴上：

第一步：构建前缀和数组

遍历权重数组w，计算前缀和prefixSum[i] = w[0] + w[1] + ... + w[i]。此时数组被映射为数轴上的连续区间：

• 下标0对应区间[0, w[0])
• 下标i对应区间[prefixSum[i-1], prefixSum[i])，区间长度为w[i]

于是"按权重随机"就等价于"在总区间[0, prefixSum[n-1])内均匀随机取一个点，看它落在哪个下标对应的区间里"。

第二步：均匀随机 + 二分定位

在[0, prefixSum[n-1])区间内随机选一个整数x，然后寻找满足x < prefixSum[i]的最小下标i，该下标即为最终解。因为区间长度恰好等于权重w[i]，所以点落到每个区间的概率与该区间长度（即权重）成正比。

对于某些下标i，所有满足prefixSum[i] - w[i] <= v < prefixSum[i]的整数v都会映射到这个下标，这也从数学上保证了"每个下标被选中的概率与下标权重成比例"。

由于prefixSum是严格递增的（所有w[i] >= 1），可以使用二分查找在 O(log n) 时间内完成定位，而不是线性扫描。

复杂度分析

• 预处理（构建前缀和）：时间复杂度 O(n)，空间复杂度 O(n)（需要保存前缀和数组）；
• pickIndex()：时间复杂度 O(log n)（一次二分查找），空间复杂度 O(1)（仅使用常数级额外变量）。

仓库源码实现解析

仓库中该题的完整实现位于 528. Random Pick with Weight.go，代码风格遵循项目一贯的命名约定（以题号作为类型与构造函数后缀，避免与其他题目类型冲突）。

package leetcode import ( "math/rand" ) // Solution528 define type Solution528 struct { prefixSum []int } // Constructor528 define func Constructor528(w []int) Solution528 { prefixSum := make([]int, len(w)) for i, e := range w { if i == 0 { prefixSum[i] = e continue } prefixSum[i] = prefixSum[i-1] + e } return Solution528{prefixSum: prefixSum} } // PickIndex define func (so *Solution528) PickIndex() int { n := rand.Intn(so.prefixSum[len(so.prefixSum)-1]) + 1 low, high := 0, len(so.prefixSum)-1 for low < high { mid := low + (high-low)>>1 if so.prefixSum[mid] == n { return mid } else if so.prefixSum[mid] < n { low = mid + 1 } else { high = mid } } return low }

关键实现细节

1. 前缀和的构建（Constructor528）

Constructor528通过一次线性遍历（O(n)）完成前缀和构建：首元素直接复制，后续元素累加前一项。最终prefixSum的最后一个元素即所有权重之和，它同时决定了随机数的取值上限。

2. 随机数的生成与区间偏移

PickIndex中的rand.Intn(total) + 1生成[1, total]区间内的整数n（total为权重总和）。与题解文档中描述的[0, prefixSum)区间相比，这里整体右移了 1，等价关系为x = n - 1，对应的映射区间相应变为(prefixSum[i-1], prefixSum[i]]，长度依然为w[i]，比例关系不受影响。

3. 二分查找的三种分支

• 当prefixSum[mid] == n时直接返回mid：此时n恰好落在区间(prefixSum[mid-1], prefixSum[mid]]的右端点上，属于下标mid的区间，直接返回是正确且省时的；
• 当prefixSum[mid] < n时，目标在右半区，low = mid + 1；
• 否则目标在左半区（含mid），high = mid。

mid := low + (high-low)>>1使用位移代替除法，并避免(low+high)/2可能出现的整数溢出问题，是二分查找的经典稳健写法。

4. 使用方式

源码末尾注释给出了标准的实例化与调用方式：

obj := Constructor(w); param_1 := obj.PickIndex();

测试验证：正确性与分支覆盖

仓库为该题提供了完整的测试用例，位于 528. Random Pick with Weight_test.go，包含两部分验证：

第一部分：基础示例验证

使用题目给出的示例权重w = [1, 3]，连续调用 6 次PickIndex并打印结果，用于人工确认输出符合权重比例（约 1/4 概率返回 0，3/4 概率返回 1）。

第二部分：边界与分支覆盖

使用五元素权重w2 = [3, 1, 1, 5, 2]（总权重 12），固定随机种子rand.Seed(1)后循环调用 2000 次，每次断言返回下标均在[0, len(w2))范围内。测试注释明确指出：多元素、多次迭代的设计是为了让二分查找充分执行每一条分支——包括low = mid + 1、high = mid以及prefixSum[mid] == n的提前返回分支。

这也呼应了仓库的整体质量保障机制：根目录下的 gotest.sh 通过go test -covermode=atomic -coverprofile=coverage.txt ./leetcode/...对全部题解执行带覆盖率统计的测试（项目 go.mod 声明 Go 1.19），确保每个题解均有测试覆盖支撑。

扩展思考

1. 为什么不用线性扫描？若每次pickIndex线性遍历查找区间，单次复杂度为 O(n)，在pickIndex最多被调用 10000 次、w.length最多 10000 的前提下最坏可达 10^8 量级；前缀和 + 二分将单次查询降为 O(log n)，是本题的标准最优解法。

2. 前缀和是必须的中间结构吗？是的。若不预计算前缀和，二分查找将无法直接在原权重数组上进行——因为二分要求查找对象具备单调性，而前缀和正是把"任意前缀区间的累计权重"这一单调序列显式化。

3. 与区间映射的联系：本解法本质上构建了一个离散化的一维概率分布，均匀随机变量经"逆变换采样"（inverse transform sampling）思路映射到离散下标，这是权重随机采样问题最通用的建模方式，可迁移至抽奖、负载均衡、采样等真实场景。

小结

528. Random Pick with Weight 的解法以"前缀和 + 二分查找"为核心：预处理阶段 O(n) 构建权重前缀和数组，查询阶段 O(log n) 完成随机下标定位，整体空间复杂度 O(n)。LeetCode-Go 仓库中的 题解文档、实现源码 与 测试代码 三者相互印证，既给出了严谨的数学映射证明，也提供了可编译、可测试、覆盖全部二分分支的完整 Go 工程实现，是学习加权随机采样算法的优质参考范本。

【免费下载链接】LeetCode-Go✅ Solutions to LeetCode by Go, 100% test coverage, runtime beats 100% | LeetCode 题解项目地址: https://gitcode.com/GitHub_Trending/le/LeetCode-Go