Random Pick with Weight

Leetcode Problem 528: Random Pick with Weight

Problem Summary

Given a 0-indexed array of positive integers w where w[i] is the weight of index i, implement a pickIndex() function that randomly selects an index such that the probability of picking index i equals w[i] / sum(w). For example, with weights [2, 6], index 0 is chosen with 25% probability and index 1 with 75% probability.

Constraints:

• The weights array has between 1 and 10,000 entries.
• Each weight is a positive integer up to 100,000.
• pickIndex() will be called at most 10,000 times.

Techniques

• Array
• Math
• Binary Search
• Prefix Sum
• Randomized

Solution

class Solution {
    private prefix: number[] = []
    private total: number = 0;

    constructor(w: number[]) {
        let sum = 0;

        for (const weight of w) {
            this.total += weight;
            this.prefix.push(this.total);
        }
    }

    pickIndex(): number {
        let left = 0
        let right = this.prefix.length - 1
        const r = Math.floor(Math.random() * this.total) + 1;

        while (left < right) {
            const mid = left + Math.floor((right - left) / 2);

            if (this.prefix[mid] < r) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }

        return left
    }
}