0528._Random_Pick_with_Weight.md

528. Random Pick with Weight

难度: Medium

刷题内容

原题连接

• https://leetcode.com/problems/random-pick-with-weight/

内容描述

Given an array w of positive integers, where w[i] describes the weight of index i, write a function pickIndex which randomly picks an index in proportion to its weight.

Note:

1 <= w.length <= 10000
1 <= w[i] <= 10^5
pickIndex will be called at most 10000 times.
Example 1:

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

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

The input is two lists: the subroutines called and their arguments. Solution's constructor has one argument, the array w. pickIndex has no arguments. Arguments are always wrapped with a list, even if there aren't any.

解题方案

思路 1 - 时间复杂度: O(len(w) * T)- 空间复杂度: O(1)******

一开始想着随机生成一个数字后，挨个减去w中的值直到小于0就取当前index返回，但是超时了，因为pickIndex的调用最多可以有10000次

时间复杂度为O(len(w) * T)，其中T为方法pickIndex的调用次数

class Solution:

    def __init__(self, w):
        """
        :type w: List[int]
        """
        self.w = w
        self.total = sum(w)
        

    def pickIndex(self):
        """
        :rtype: int
        """
        position = random.random() * self.total
        for i in range(len(self.w)):
            position -= self.w[i]
            if position <= 0: # 这里如果用小于0的话，那么当正好取到0的时候永远没有答案
                return i

思路 2 - 时间复杂度: O(lg(len(w)) * T)- 空间复杂度: O(N)******

于是想着用一个前缀和数组来辅助，直接二分查找就可以了

beats 46.38%

class Solution:

    def __init__(self, w):
        """
        :type w: List[int]
        """
        self.pref = [w[0]] * len(w)
        for i in range(1, len(w)):
            self.pref[i] = self.pref[i-1] + w[i]
        self.total = sum(w)
        

    def pickIndex(self):
        """
        :rtype: int
        """
        seed = random.randint(1, self.total)
        l, r = 0, len(self.pref) - 1
        while l <= r:
            mid = l + ((r-l) >> 1)
            if seed < self.pref[mid]:
                r = mid - 1
            elif seed > self.pref[mid]:
                l = mid + 1
            else:
                return mid
        return l