Random Pick with Weight — Prefix Sums and Probabilistic Sampling

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(int[] w) — Initializes the object with the array w.
• int pickIndex() — Returns a random index in the range [0, w.length - 1] with probability proportional to w[i] / sum(w).

Constraints:

• 1 <= w.length <= 10^4
• 1 <= w[i] <= 10^5
• pickIndex will be called at most 10^4 times.

Example 1:
Input:  w = [1], pickIndex() called
Output: 0
Explanation: Only one choice; always returns 0.

Example 2:
Input:  w = [1, 3], pickIndex() called multiple times
Output: Approximately 0 about 25% and 1 about 75% of the time.
Explanation: Total weight = 4. P(0) = 1/4 = 25%, P(1) = 3/4 = 75%.

Example 3:
Input:  w = [1, 3, 2], pickIndex()
Output: Index 0 with prob 1/6, index 1 with prob 3/6 = 1/2, index 2 with prob 2/6 = 1/3.

Why This Problem Matters#

Random Pick with Weight is the algorithmic foundation of weighted random sampling — a technique used in machine learning (stochastic gradient descent sampling), A/B testing (traffic splitting by weight), recommendation systems (diversity-weighted selection), and load balancing (route traffic proportional to server capacity). Google and Meta ask this problem because it sits at the intersection of probability, data structures, and practical systems design.

The problem appears simple — "return a random index with probability proportional to its weight" — but the implementation requires connecting three ideas: (1) prefix sums to convert weights into intervals, (2) random number generation within the total weight range, and (3) binary search to efficiently find which interval the random number falls into.

Many candidates reach for the naive approach: expand the weights into a flat array (e.g., w = [1,3] becomes [0, 1, 1, 1]) and pick a random index. This works correctly but uses O(sum(w)) space — potentially hundreds of millions of entries. The prefix sum + binary search approach uses O(n) space and O(log n) per pick, which is the expected answer in an interview.

This problem also reinforces the canonical "binary search on prefix sums" pattern that appears in many other problems: finding the right bucket in a probability distribution, sampling from a cumulative distribution function, and even some range query problems. Mastering it here pays dividends across a broad problem family.

The Core Insight#

Convert weights to a probability distribution using prefix sums.

Prefix sum array prefix[i] = sum(w[0..i]). After building this array:

• prefix = [1, 4, 6] for w = [1, 3, 2]
• Weight 1 corresponds to the range [1, 1]
• Weight 3 corresponds to the range [2, 4]
• Weight 2 corresponds to the range [5, 6]

Pick a random integer in [1, total_weight].

A uniformly random integer in [1, total] falls in range [prefix[i-1]+1, prefix[i]] with probability w[i] / total, which is exactly the desired probability for index i.

Find the range using binary search.

Given random r, find the smallest index i such that prefix[i] >= r. This is bisect_left(prefix, r) in Python or a standard lower-bound binary search. This index is the sampled index.

The combination: build prefix sums in __init__ (O(n)), then answer each pickIndex call with a single random number generation and a binary search (O(log n)).

Visual Dry Run#

Input: w = [1, 3, 2]

Build prefix sums: prefix = [1, 4, 6]

Total weight = 6.

Sampling space:

Range [1,1]:  maps to index 0 (weight 1, prob = 1/6)
Range [2,4]:  maps to index 1 (weight 3, prob = 3/6)
Range [5,6]:  maps to index 2 (weight 2, prob = 2/6)

pickIndex() — random r = 3: Binary search for smallest prefix ≥ 3 in [1, 4, 6].

• prefix[0] = 1 < 3 → not here
• prefix[1] = 4 >= 3 → found at index 1

Return index 1.

pickIndex() — random r = 6: Binary search: prefix[2] = 6 >= 6 → found at index 2. Return 2.

pickIndex() — random r = 1: Binary search: prefix[0] = 1 >= 1 → found at index 0. Return 0.

Solution (Optimal)#

import random
import bisect
 
class Solution:
    def __init__(self, w: list[int]):
        # Build prefix sum array
        self.prefix = []
        total = 0
        for weight in w:
            total += weight
            self.prefix.append(total)
        self.total = total
 
    def pickIndex(self) -> int:
        # Random integer in [1, total]
        r = random.randint(1, self.total)
        # Find smallest index i such that prefix[i] >= r
        return bisect.bisect_left(self.prefix, r)
 
 
# Manual binary search (without bisect module)
class Solution_manual:
    def __init__(self, w: list[int]):
        self.prefix = []
        running = 0
        for weight in w:
            running += weight
            self.prefix.append(running)
        self.total = running
 
    def pickIndex(self) -> int:
        r = random.randint(1, self.total)
        lo, hi = 0, len(self.prefix) - 1
        while lo < hi:
            mid = (lo + hi) // 2
            if self.prefix[mid] < r:
                lo = mid + 1
            else:
                hi = mid
        return lo

class Solution {
    constructor(w) {
        this.prefix = [];
        let total = 0;
        for (const weight of w) {
            total += weight;
            this.prefix.push(total);
        }
        this.total = total;
    }
 
    pickIndex() {
        // Random float in [0, 1) scaled to [1, total]
        const r = Math.floor(Math.random() * this.total) + 1;
 
        // Binary search: find smallest index with prefix[i] >= r
        let lo = 0, hi = this.prefix.length - 1;
        while (lo < hi) {
            const mid = Math.floor((lo + hi) / 2);
            if (this.prefix[mid] < r) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        return lo;
    }
}

Complexity Analysis

Compared to the naive flat-array expansion: O(sum(w)) space vs. O(n) space, and O(1) pick vs. O(log n) pick. For large weights, the prefix sum approach is far superior in memory.

Common Mistakes#

• Using random.random() and scaling to [0, total) with math.floor. Floating-point arithmetic can cause edge cases near integer boundaries. Using random.randint(1, total) (integer arithmetic) is safer and cleaner.
• Using bisect_right instead of bisect_left. bisect_left returns the leftmost position where r could be inserted to keep the array sorted — i.e., the first index where prefix[i] >= r. This is correct. bisect_right returns the position after all existing r values, which would incorrectly skip over exact matches.
• Forgetting to initialize with 1-indexed random (not 0-indexed). If you generate r in [0, total) and compare with prefix[i], you need to use strict less-than comparison. Using [1, total] and prefix[i] >= r (non-strict) is cleaner and less error-prone.
• Not handling the single-element case. When w = [5], prefix = [5], total = 5, and r is always 5. bisect_left([5], 5) = 0. Correct — returns index 0. No special case needed.
• Storing weights instead of prefix sums. The whole point is to precompute prefix sums for O(log n) lookup. Storing raw weights requires O(n) lookup per pick.

Follow-up Questions#

What is the difference between random.randint(1, total) and random.randint(0, total-1)? They are equivalent in probability but map differently to the prefix sum array. Using [1, total] maps range [prefix[i-1]+1, prefix[i]] to index i. Using [0, total-1] maps range [prefix[i-1], prefix[i]-1] to index i — requires a slightly different binary search condition.

What if pickIndex is called billions of times — is O(log n) fast enough? Yes. For n = 10^4, log2(10^4) ≈ 14 operations per pick. At billions of calls, the dominant cost is random number generation, not the binary search.

How would you sample without replacement (each index can be picked at most once)? After each pick, set w[picked_index] = 0 and rebuild prefix sums. More efficiently, use the Fisher-Yates shuffle on a weighted array and iterate through the shuffled order.

Can you use random.choices (Python) instead of the manual implementation? Yes. random.choices(range(len(w)), weights=w, k=1)[0] handles weighted sampling internally. In an interview, you should implement the algorithm manually to demonstrate understanding.

What if new weights are added dynamically (insert operations)? Rebuild the prefix sum array on each insert: O(n) per insert. For truly dynamic weighted sampling, use a Fenwick tree (Binary Indexed Tree) to support O(log n) weight updates and O(log n) weighted sampling.

Key Takeaways#

• LC 528 Random Pick with Weight maps weighted sampling to a uniform [0, total) draw plus binary search.
• Build a prefix-sum array of weights once in O(n); each pick is O(log n) via bisect_left (Python) or manual binary search (JavaScript).
• The probability of picking index i is w[i] / total, exactly as required.
• Use random() * total (continuous) or randint(1, total) (discrete) — both work, just be consistent with bisect bounds.
• Avoid building a flattened repeat array — it can be O(sum_of_weights) memory and blows up for large weights.
• Time per pick O(log n), space O(n) — optimal for static weights.
• The same prefix-sum + binary-search pattern powers ad auctions, A/B test bucket assignment, and weighted load balancing.

Related Problems#

• LC 528 — Random Pick with Weight: This problem.
• LC 398 — Random Pick Index: Pick a random index from a list where the target value appears — reservoir sampling.
• LC 710 — Random Pick with Blacklist: Uniform random pick from [0, n-1] excluding blacklisted indices.
• LC 470 — Implement Rand10() Using Rand7(): Probability transformation using a known random source.
• LC 384 — Shuffle an Array: Uniform random shuffle — uses the Fisher-Yates algorithm.
• LC 380 — Insert Delete GetRandom O(1): Uniform random pick from a dynamic set.