按权重随机选择

难度: Medium

给定一个正整数数组 w ，其中 w[i] 代表下标 i 的权重（下标从 0 开始），请写一个函数 pickIndex ，它可以随机地获取下标 i，选取下标 i 的概率与 w[i] 成正比。

例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25 （即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。

也就是说，选取下标 i 的概率为 w[i] / sum(w) 。

示例 1：

输入：
inputs = ["Solution","pickIndex"]
inputs = [[[1]],[]]
输出：
[null,0]
解释：
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。

示例 2：

输入：
inputs = ["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
inputs = [[[1,3]],[],[],[],[],[]]
输出：
[null,1,1,1,1,0]
解释：
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // 返回 1，返回下标 1，返回该下标概率为 3/4 。
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 0，返回下标 0，返回该下标概率为 1/4 。

由于这是一个随机问题，允许多个答案，因此下列输出都可以被认为是正确的:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
诸若此类。

提示：

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

注意：本题与主站 528 题相同： https://leetcode-cn.com/problems/random-pick-with-weight/

Submission

运行时间: 99 ms

内存: 20.9 MB

class Solution:

    def __init__(self, w: List[int]):
        self.prob = [0]
        _sum = sum(w)
        temp = 0
        for i in w:
            temp += i
            self.prob.append(temp / _sum)
    
    def pickIndex(self) -> int:
        return bisect.bisect_left(self.prob, random.random()) - 1



# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()

Explain

此题解采用了前缀和和二分搜索的方法来解决问题。首先，将输入权重数组转换为一个概率数组，该数组的每个元素表示从开始到当前位置的权重和占总权重的比例。具体操作是，在初始化函数中计算总权重，并构建一个概率数组，其中每个元素是到当前元素为止的权重和除以总权重。在pickIndex函数中，使用random.random()生成一个[0, 1)之间的随机数，然后通过二分搜索（bisect_left）找到这个随机数应该插入概率数组的位置，由于概率数组是累积的，这样可以保证按权重随机选择下标。

时间复杂度: O(log n)

空间复杂度: O(n)

class Solution:

    def __init__(self, w: List[int]):
        # 初始化概率数组
        self.prob = [0]
        # 计算权重的总和
        _sum = sum(w)
        temp = 0
        for i in w:
            temp += i
            # 构建概率分布数组
            self.prob.append(temp / _sum)
    
    def pickIndex(self) -> int:
        # 生成随机数并通过二分查找返回对应的下标
        return bisect.bisect_left(self.prob, random.random()) - 1

# 使用示例
# obj = Solution(w)
# param_1 = obj.pickIndex()

Explore

在初始化概率数组时，第一个元素设置为0是为了配合二分搜索算法的使用。这种设置使得概率数组成为一个从0开始，逐渐增加直到1的数组。由于random.random()生成的随机数范围是[0, 1)，第一个元素为0可以确保即使随机数非常小（接近0），二分搜索也能正确找到第一个大于或等于这个随机数的位置。这保证了在边界情况下，返回的索引能够正确反映权重数组中第一个元素的权重。

是的，这种方法确保了概率分布正确地反映了权重的累积特性。通过逐个添加权重到`temp`并将其除以总权重来更新`self.prob`，我们逐步构建了一个累积概率数组。这个累积概率数组表示从第一个元素到当前元素的权重和占总权重的比例。因此，每个元素在概率数组中的位置正比于其权重，从而确保了按权重的随机选择是公正的。

在`pickIndex`函数中，使用`random.random()`生成的随机数代表在[0, 1)区间内的一个随机实数。这个随机数用来模拟均匀分布的随机事件，其用途是从按权重设置的概率分布中选择索引。通过将这个随机数与累积的概率数组进行比较，我们可以通过二分搜索找到第一个大于或等于这个随机数的概率值的位置。这使得返回的索引的概率与其相应权重成正比，从而确保了按权重的随机选择。换言之，权重越大的元素，其累积概率区间也越大，因此被随机选中的概率也越高。