requirements

You are given an array w of positive integers with subscripts starting at 0, where w[I] represents the weight of the ith subscript.

Implement a pickIndex function that randomly selects and returns a subscript from the range [0, w.length-1] (including 0 and w.Length-1). The probability of subscript I is w[I] / sum(w).

For example, for w = [1, 3], the probability of picking subscript 0 is 1 / (1 + 3) = 0.25 (i.e., 25%), while the probability of picking subscript 1 is 3 / (1 + 3) = 0.75 (i.e., 75%).

Example 1:

Input: ["Solution","pickIndex"] [[[1]] output: [null,0] explanation: Solution Solution = new Solution([1]); solution.pickIndex(); // Return 0, since there is only one element in the array, the only option is to return the index 0.Copy the code

Example 2:

Input: [" Solution ", "pickIndex", "pickIndex", "pickIndex", "pickIndex", "pickIndex"] [[[1, 3]], [], [], [], [], []] output: [null,1,1,1,1,0] [null,1,1,1, 0] solution.pickIndex(); // Returns 1, returns subscript 1 with probability 3/4. solution.pickIndex(); // Return 1 solution.pickindex (); // Return 1 solution.pickindex (); // Return 1 solution.pickindex (); // Returns 0, returns subscript 0 with probability 1/4. Since this is a random question, allowing for multiple answers, the following output can be considered correct: [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]... So on.Copy the code

The core code

import random
import bisect
class Solution:

    def __init__(self, w: List[int]):
        self.w = w
        self.sw = []
        w_sum = 0
        for ww in self.w:
            w_sum += ww
            self.sw.append(w_sum)

    def pickIndex(self) -> int:
        return bisect.bisect_left(self.sw,random.randint(1,self.sw[-1]))
Copy the code

Answer: We used the bisect method. If you’re not familiar with it, the Python Bisect module randomly picks position I, and the probability of picking position I is proportional to w[I]. It’s not related to the probability calculation of the output in this case. Random finally output coincidence passed this problem, learn bisect binary search method on the line.