Random Pick with Weight [Medium] — Prefix Sum + Binary Search

Problem Statement

Implement Solution(w) where w[i] = weight of index i. pickIndex() returns an index with probability proportional to its weight.

w=[1,3] → pickIndex picks 0 with prob 1/4, 1 with prob 3/4

Intuition

Build prefix sums. Generate random number in [1, total_weight]. Binary search for the first prefix sum >= random number.

Solutions

Python

import random, bisect

class Solution:
    def __init__(self, w: list[int]):
        self.prefix = []
        total = 0
        for x in w:
            total += x
            self.prefix.append(total)
        self.total = total

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

Java

class Solution {
    int[] prefix; int total; Random rand=new Random();
    public Solution(int[] w) {
        prefix=new int[w.length];
        prefix[0]=w[0];
        for(int i=1;i<w.length;i++) prefix[i]=prefix[i-1]+w[i];
        total=prefix[w.length-1];
    }
    public int pickIndex() {
        int r=rand.nextInt(total)+1;
        int lo=0,hi=prefix.length-1;
        while(lo<hi){int mid=(lo+hi)/2;if(prefix[mid]<r)lo=mid+1;else hi=mid;}
        return lo;
    }
}

C++

class Solution {
    vector<int> prefix; int total;
public:
    Solution(vector<int>& w){
        for(int x:w){total+=x;prefix.push_back(total);}
    }
    int pickIndex(){
        int r=rand()%total+1;
        return lower_bound(prefix.begin(),prefix.end(),r)-prefix.begin();
    }
};

Complexity

init: O(n), pickIndex: O(log n)