528. 按权重随机选择

Difficulty
Medium
Tags
二分查找
URL
https://leetcode-cn.com/problems/random-pick-with-weight/
Star
给你一个 下标从 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 <= 104
  • 1 <= w[i] <= 105
  • pickIndex 将被调用不超过 104 次
通过次数32,625提交次数68,226

二分法

思路
根据原始数组值建立一个权重数组。
notion imagenotion image
根据权重数组范围从[1, 11]随机挑选一个数字，看它落在哪个区间，然后该区间的下标即为选中的数字。思想即为原数组数字大的概率大，那就让该空间也大一点，类似于转盘扔飞镖，概率越大的对应面积也越大。
题解
import random
class Solution:

    def __init__(self, w: List[int]):
        self.arr = w
        # 权重数组，抽取的随机数落在 ( weight[i-1] 于 weight[i] ]之间则选择 i
        self.weight = [num for num in self.arr]
        for i in range(1, len(self.weight)):
            self.weight[i] = self.weight[i-1] + self.arr[i]


    def pickIndex(self) -> int:
        randVal = random.randint(1, self.weight[-1])
        # 用二分法查找 该随机数在权重矩阵中的范围
        # 即查找 randVal 在weight 大于等于 该值的最小值
        left, right = 0, len(self.weight)-1
        while left <= right:
            mid = left + (right - left) // 2
            if self.weight[mid] < randVal:
                left = mid + 1
            else:
                right = mid - 1

        return left

# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()
注：random.randint(min, max) 取值范围为 [min, max] 左右均为闭区间