Random Pick with Weight Python Solution

lcpy gen -n 528   # by problem number
lcpy gen -s random_pick_with_weight   # by problem name

Input
["Solution","pickIndex"]
[[[1]],[]]
Output
[null,0]

Explanation
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. The only option is to return 0 since there is only one element in w.

Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output
[null,1,1,1,1,0]

Explanation
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // return 1. It is returning the second element (index = 1) that has a probability of 3/4.
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 0. It is returning the first element (index = 0) that has a probability of 1/4.

Since this is a randomization problem, multiple answers are allowed.

import bisect
import random


class Solution:
    # Build prefix sums of weights. pick_index draws r in [1, total],
    # binary search for first prefix >= r. Each index i chosen with prob w[i]/sum.
    # Time: O(n) init, O(log n) pick_index
    # Space: O(n)
    def __init__(self, w: list[int]) -> None:
        self.prefix: list[int] = []
        running = 0
        for weight in w:
            running += weight
            self.prefix.append(running)
        self.total = running

    def pick_index(self) -> int:
        r = random.randint(1, self.total)
        return bisect.bisect_left(self.prefix, r)