按权重随机选择

难度: Medium

给你一个 下标从 0 开始 的正整数数组 w ，其中 w[i] 代表第 i 个下标的权重。

请你实现一个函数 pickIndex ，它可以 随机地 从范围 [0, w.length - 1] 内（含 0 和 w.length - 1）选出并返回一个下标。选取下标 i 的概率为 w[i] / sum(w) 。

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

示例 1：

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

示例 2：

输入：
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[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 <= 10⁴
1 <= w[i] <= 10⁵
pickIndex 将被调用不超过 10⁴ 次

Submission

运行时间: 136 ms

内存: 21.0 MB

class Solution:

    def __init__(self, w: List[int]):
        self.pre = list(accumulate(w))
        self.total = sum(w)

    def pickIndex(self) -> int:
        x = random.randint(1, self.total)
        return bisect_left(self.pre, x)

Explain

此题解使用前缀和数组和二分查找算法来实现按权重随机选择的功能。首先，在初始化函数中，计算输入权重数组的前缀和，并存储总权重。在pickIndex函数中，随机生成一个介于1到总权重之间的整数，然后使用二分查找在前缀和数组中找到这个随机数应该插入的位置，该位置即为所选的下标。

时间复杂度: O(log n)

空间复杂度: O(n)


class Solution:

    def __init__(self, w: List[int]):
        # 计算权重的前缀和，便于后续的快速查找
        self.pre = list(accumulate(w))
        # 计算总权重，用于生成随机数
        self.total = sum(w)

    def pickIndex(self) -> int:
        # 生成一个[1, total]范围内的随机数
        x = random.randint(1, self.total)
        # 使用二分查找法找到随机数x对应的下标
        return bisect_left(self.pre, x)

Explore

使用`accumulate(w)`计算前缀和在时间复杂度上是O(n)，其中n是数组`w`的长度。因此，对于非常大的数组，这种计算确实会消耗相对较多的时间。然而，这种计算只在初始化阶段进行一次，因此总体对性能的影响相对较小。如果关注内存使用，前缀和数组的空间复杂度同样是O(n)，这可能在极大的数据集上造成内存压力。

在`pickIndex`函数中生成随机数的范围是从1到总权重`self.total`，这是因为权重代表了选择每个元素的相对可能性，且前缀和数组的最小值从1开始（假设所有输入权重至少为1）。如果从0开始生成随机数，则可能导致在查找时出现数组索引越界的问题，或者使得第一个元素的选中概率不均匀。

当使用`bisect_left`进行二分查找时，如果随机数正好等于前缀和数组中的某个值，`bisect_left`会返回该值在数组中的最左侧位置。这是因为`bisect_left`的目的是找到插入点，以保持数组的有序性。因此，即使有多个相同的值，它也会返回第一个这样的值的位置。这确保了随机选择权重的准确性。

前缀和数组是单调递增的，这使得二分查找成为寻找随机数位置的理想选择，因为二分查找的时间复杂度是O(log n)，非常高效。相比如线性搜索（O(n)时间复杂度），二分查找显著减少了查找时间，特别是在数组非常大时。此外，二分查找直接跳到中间值比较，减少了不必要的比较，是处理有序数据的最有效方法之一。