Random Pick with Weight

Problem Statement

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the i-th index. Implement the Solution class:

• Solution(w) — Initialize with the weights array.
• pickIndex() — Return a random index i such that the probability of picking index i is w[i] / sum(w).

Example:

Input: w = [1, 3]
Explanation: 
- Index 0 has weight 1, index 1 has weight 3. Total = 4.
- pickIndex() should return 0 with probability 1/4 and 1 with probability 3/4.

Constraints:

• 1 <= w.length <= 10^4
• 1 <= w[i] <= 10^5
• At most 10^4 calls to pickIndex.

Difficulty: Medium (LC 528)

---

Pattern Recognition

• What is the input structure? An array of weights defining a probability distribution.
• What are we optimizing? O(log n) random sampling from a weighted distribution.
• What pattern does this fit? Prefix Sum + Binary Search. Build a cumulative weight array. Generate a random number in the total range, then binary search for the corresponding index.
• Why does this pattern fit? The prefix sum transforms the weights into non-overlapping ranges on a number line. Each index "owns" a range proportional to its weight. A uniform random number in [1, total] lands in exactly one range, and binary search finds that range in O(log n).

---

Approach

Brute Force — Expand and Pick

Create an array where index i appears w[i] times. Pick a random element.

expanded = []
for i, weight in enumerate(w):
    expanded += [i] * weight
return random.choice(expanded)

Time: O(sum(w)) for initialization, O(1) for pick. Space: O(sum(w)) — can be enormous (up to 10^9). Why insufficient: Space is proportional to the sum of weights, not the number of indices. With weights up to 10^5 and 10^4 indices, the expanded array could have 10^9 elements.

Optimal — Prefix Sum + Binary Search

Key Insight: Build a prefix sum array where prefix[i] = w[0] + w[1] + ... + w[i]. This creates non-overlapping ranges: index 0 owns [1, w[0]], index 1 owns [w[0]+1, w[0]+w[1]], etc. Generate a random number in [1, total] and binary search for the first prefix sum >= that number.

Step-by-step:

1. Init: Build prefix sum array. prefix[i] = prefix[i-1] + w[i].
2. pickIndex:
   • Generate random integer target in [1, prefix[-1]].
   • Binary search for the leftmost index where prefix[index] >= target.
   • Return that index.

class Solution:
    init(w):
        prefix = cumulative sum of w
    
    pickIndex():
        target = random(1, prefix[-1])
        return bisect_left(prefix, target)

Time: O(n) init, O(log n) per pick. Space: O(n).

---

Walkthrough

w = [2, 5, 3, 4]

Step 1 — Build prefix sum:

Index	Weight	Prefix Sum	Owns range
0	2	2	[1, 2]
1	5	7	[3, 7]
2	3	10	[8, 10]
3	4	14	[11, 14]

Total = 14.

Step 2 — pickIndex examples:

Random target	bisect_left(prefix, target)	Result
1	prefix = [2,7,10,14], bisect_left for 1 → 0	Index 0
2	bisect_left for 2 → 0	Index 0
3	bisect_left for 3 → 1	Index 1
7	bisect_left for 7 → 1	Index 1
8	bisect_left for 8 → 2	Index 2
11	bisect_left for 11 → 3	Index 3
14	bisect_left for 14 → 3	Index 3

Probability of each index: 2/14, 5/14, 3/14, 4/14 — proportional to weights.

---

Implementation

TypeScript

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

    constructor(w: number[]) {
        // Build prefix sum array
        this.prefix = [];
        let running = 0;
        for (const weight of w) {
            running += weight;
            this.prefix.push(running);
        }
        this.total = running;
    }

    pickIndex(): number {
        // Random target in [1, total]
        const target = Math.floor(Math.random() * this.total) + 1;

        // Binary search: find the leftmost index where prefix[index] >= target
        let lo = 0, hi = this.prefix.length - 1;
        while (lo < hi) {
            const mid = (lo + hi) >> 1;
            if (this.prefix[mid] < target) lo = mid + 1;
            else hi = mid;
        }
        return lo;
    }
}

C++

class Solution {
    vector<int> prefix;
    int total;

public:
    Solution(vector<int>& w) {
        // Build prefix sum array
        int running = 0;
        for (int weight : w) {
            running += weight;
            prefix.push_back(running);
        }
        total = running;
    }

    int pickIndex() {
        // Random target in [1, total]
        int target = rand() % total + 1;

        // Find the leftmost index where prefix[index] >= target
        return (int)(lower_bound(prefix.begin(), prefix.end(), target) - prefix.begin());
    }
};

Alternative — using itertools.accumulate:

TypeScript

class SolutionClean {
    private prefix: number[];
    private total: number;

    constructor(w: number[]) {
        // Build prefix sum using reduce
        this.prefix = w.reduce<number[]>((acc, val) => {
            acc.push((acc.length > 0 ? acc[acc.length - 1] : 0) + val);
            return acc;
        }, []);
        this.total = this.prefix[this.prefix.length - 1];
    }

    pickIndex(): number {
        const target = Math.floor(Math.random() * this.total) + 1;
        let lo = 0, hi = this.prefix.length - 1;
        while (lo < hi) {
            const mid = (lo + hi) >> 1;
            if (this.prefix[mid] < target) lo = mid + 1;
            else hi = mid;
        }
        return lo;
    }
}

C++

class SolutionClean {
    vector<int> prefix;
    int total;

public:
    SolutionClean(vector<int>& w) {
        partial_sum(w.begin(), w.end(), back_inserter(prefix));
        total = prefix.back();
    }

    int pickIndex() {
        int target = rand() % total + 1;
        return (int)(lower_bound(prefix.begin(), prefix.end(), target) - prefix.begin());
    }
};

Manual binary search version (shows understanding):

TypeScript

class SolutionManual {
    private prefix: number[];
    private total: number;

    constructor(w: number[]) {
        this.prefix = [];
        let running = 0;
        for (const weight of w) {
            running += weight;
            this.prefix.push(running);
        }
        this.total = running;
    }

    pickIndex(): number {
        const target = Math.floor(Math.random() * this.total) + 1;

        let lo = 0, hi = this.prefix.length - 1;
        while (lo < hi) {
            const mid = (lo + hi) >> 1;
            if (this.prefix[mid] < target) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return lo;
    }
}

C++

class SolutionManual {
    vector<int> prefix;
    int total;

public:
    SolutionManual(vector<int>& w) {
        partial_sum(w.begin(), w.end(), back_inserter(prefix));
        total = prefix.back();
    }

    int pickIndex() {
        int target = rand() % total + 1;

        int lo = 0, hi = (int)prefix.size() - 1;
        while (lo < hi) {
            int mid = (lo + hi) / 2;
            if (prefix[mid] < target) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return lo;
    }
};

Watch out:

1. Random range: [0, total) vs [1, total]. If you use random.random() * total, you get a float in [0, total). Then use bisect_right instead of bisect_left, because prefix[i] is the upper bound of index i's range. Alternatively, use random.randint(1, total) for integers and bisect_left — this is cleaner.
2. Using bisect_right instead of bisect_left. With integer targets in [1, total]: bisect_left(prefix, target) finds the first prefix >= target, which is the correct index. bisect_right would skip to the next index when there is an exact match, giving wrong probabilities.
3. Off-by-one in prefix sum construction. The prefix sum must be inclusive: prefix[0] = w[0], not 0. If you prepend a 0, adjust the binary search accordingly.

---

Complexity Analysis

	Time (init)	Time (pickIndex)	Space
Brute Force (Expand)	O(sum(w))	O(1)	O(sum(w))
Prefix Sum + Binary Search	O(n)	O(log n)	O(n)

Bottleneck: The O(log n) per pick is optimal for this approach. There are more advanced techniques (alias method) that achieve O(1) per pick with O(n) preprocessing, but they are complex and rarely expected in interviews.

---

Edge Cases

1. Single element — Always return index 0. Binary search trivially works.
2. All equal weights — Uniform distribution. Each index has probability 1/n.
3. One very large weight, rest small — That index dominates. Most random picks return it.
4. Weight of 1 for all — Same as uniform random choice over indices.
5. Large total weight (10^9) — Python handles large integers natively. In other languages, check for overflow in prefix sum.
6. Many pickIndex calls — Each call is O(log n), independent. No state changes between calls.

---

Related Problems

• LC 398 — Random Pick Index — Reservoir sampling: pick random index of a target value in a stream.
• LC 382 — Linked List Random Node — Reservoir sampling on a linked list.
• LC 497 — Random Point in Non-overlapping Rectangles — Same prefix sum + binary search, but with 2D rectangles.
• LC 710 — Random Pick with Blacklist — Uniform random pick excluding certain values. Uses remapping.

---

What is Expected at Each Level

Junior: Understand the expanded array approach and why it is wasteful. With guidance, arrive at prefix sum + binary search. May struggle with the exact binary search variant (bisect_left vs bisect_right).

Mid-level: Immediately see prefix sum + binary search. Write clean code using bisect_left. Explain why the probability is correct. Handle edge cases. Discuss the random number range carefully.

Senior: Solve quickly. Discuss the Alias Method for O(1) per pick (Vose's algorithm). Explain the tradeoff: O(n) preprocessing for O(1) picks. Discuss applications: weighted load balancing, sampling in ML, Monte Carlo simulations. Mention reservoir sampling for streaming scenarios where weights arrive online.