Random Pick with Weight: Prefix Sums + Binary Search

LeetCode 528, Random Pick with Weight, asks you to pick a random index from an array where each index has a different probability of being chosen based on its weight. Index i should be picked with probability w[i] / sum(w). This is a classic combination of prefix sums and binary search, and it shows up in everything from load balancers to Monte Carlo simulations.

The problem

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

• __init__(self, w): Initializes the object with the weight array
• pickIndex(self) -> int: Returns a random index i where the probability of picking i is w[i] / sum(w)

w = [1, 3, 2, 4]

pickIndex()  # returns 0 with probability 1/10 (10%)
pickIndex()  # returns 1 with probability 3/10 (30%)
pickIndex()  # returns 2 with probability 2/10 (20%)
pickIndex()  # returns 3 with probability 4/10 (40%)

Think of it like a spinner divided into sections of different sizes. Each section corresponds to an index, and its size is proportional to its weight. You spin the spinner and report which section it lands on.

The key insight

The trick is to turn the weight array into a number line using prefix sums, then use binary search to find which segment a random number falls into.

If w = [1, 3, 2, 4], the prefix sums are [1, 4, 6, 10]. Now generate a random integer in [1, 10]. If you get 1, that belongs to index 0 (range 1 to 1). If you get 2, 3, or 4, that belongs to index 1 (range 2 to 4). If you get 5 or 6, that belongs to index 2 (range 5 to 6). If you get 7 through 10, that belongs to index 3 (range 7 to 10).

Each index "owns" exactly w[i] numbers out of the total, so the probabilities are correct by construction.

The solution

import random
import bisect

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

    def pickIndex(self):
        target = random.randint(1, self.total)
        return bisect.bisect_left(self.prefix, target)

The constructor builds the prefix sum array in O(n). Each call to pickIndex generates a random number in [1, total] and uses bisect_left to find the first prefix sum that is greater than or equal to the target. That index is our answer.

Why bisect_left? Because we want the leftmost position where prefix[i] >= target. If target = 4 and prefix = [1, 4, 6, 10], bisect_left returns index 1, which is correct since index 1 owns the range 2 to 4.

Step-by-step walkthrough

Let's trace through the full algorithm with w = [1, 3, 2, 4]. First we build the prefix sums, then we simulate a pickIndex() call with a random number of 5.

Complexity analysis

The prefix sum approach is the clear winner. The expansion approach creates an array with sum(w) elements (repeating each index by its weight count), which blows up in memory for large weights. Prefix sums keep space at O(n) regardless of the weight values, and binary search keeps each pick at O(log n).

The building blocks

Prefix sums

Prefix sums turn an array of values into an array of running totals. This lets you answer "what is the sum of elements from index i to j?" in O(1) after O(n) preprocessing. Here, we use prefix sums differently: they define the boundaries of weighted segments on a number line. The same technique shows up in range sum queries, subarray sum problems, and anywhere you need cumulative totals.

Binary search on a sorted boundary

Once you have the prefix sum array, it is sorted by definition (since all weights are positive). Binary search finds which segment a random number falls into. This is the same "find the leftmost element satisfying a condition" pattern that appears in problems like Search Insert Position and Find First and Last Position. The condition here is prefix[i] >= target.

Edge cases

• Single element: If w = [5], the prefix array is [5], and every random number from 1 to 5 maps to index 0. The single element is always returned.
• All equal weights: If w = [2, 2, 2], each index gets a 1/3 probability. The prefix array is [2, 4, 6], dividing the number line into three equal segments.
• Very large weights: Since we only store prefix sums (not expanded arrays), even w = [1000000, 1000000] only needs 2 entries. The approach scales with array length, not weight magnitude.
• Weight of 1 for one element: That index gets a tiny slice but is never excluded. For w = [1, 999], index 0 still has a 0.1% chance.
• Multiple pickIndex calls: Each call is independent. The random number generation is fresh each time, and no state carries over between calls.

From understanding to recall

This problem ties together two fundamental patterns: prefix sums and binary search on a sorted structure. Understanding why it works is one thing, but reproducing it cleanly in an interview requires that both pieces are automatic. You need to build prefix sums without hesitating over off-by-one errors, and you need to reach for bisect_left (or write the equivalent loop) without second-guessing the comparison direction.

Spaced repetition locks in both patterns. Practice building the prefix array and writing the binary search from scratch at increasing intervals. After a few rounds, the combination feels like a single move rather than two separate ideas stitched together.

Related posts

• Binary Search - The foundational binary search template that powers the lookup in this problem
• Search Insert Position - Binary search for the correct insertion point, the same pattern used to find the right prefix sum bucket
• Random Pick Index - A related randomized selection problem using reservoir sampling instead of prefix sums