How do you solve Random Pick with Weight?

Precompute prefix sums of the weights during initialization. For each pickIndex call, generate a random integer target between 1 and total weight inclusive. Use binary search on the prefix sums to find the index where prefix_sum >= target. Return this index as the weighted random pick.

What pattern does Random Pick with Weight use?

The Pattern: Binary Search on Prefix Sums for Weighted Random Selection The Heuristic: Transform weighted probabilities into prefix sums and use binary search to efficiently map a random target to its corresponding weight interval, enabling O(log n) selection time.

Random Pick with Weight — Modified Binary Search

Q: What is the time complexity of Random Pick with Weight?

O(n) for initialization, O(log n) per pickIndex call. Computing prefix sums requires a single pass through the weights (O(n)). Each pickIndex call uses binary search on the prefix sums array, which takes O(log n) time.

Problem

Given a list of positive integer weights, implement a function that randomly picks an index with probability proportional to its weight.

1 ≤ w.length ≤ 10⁴
1 ≤ w[i] ≤ 10⁵
pickIndex will be called at most 10⁴ times

Example

Input: w = [1, 3]

Output: 0 with probability 1/4, 1 with probability 3/4

The total weight is 4. The algorithm first computes prefix sums: [1, 4]. When pickIndex is called, it generates a random integer target between 1 and 4 inclusive. If target is 1, it returns index 0; if target is 2, 3, or 4, it returns index 1. This ensures the probability of returning index i is proportional to w[i]. The critical moment is using binary search on the prefix sums to efficiently find the correct index for the random target.

Approach

Straightforward Solution

A naive approach would generate a random number and scan through the weights cumulatively to find the corresponding index, resulting in O(n) time per pickIndex call, which is inefficient for large inputs or many calls.

Core Observation

The probability of picking an index i is w[i] divided by the total sum of weights. This can be modeled as selecting a random number in the range [1, total_weight] and finding which prefix sum interval it falls into.

Path to Optimal

Preview

By precomputing prefix sums, the problem reduces to searching for the smallest prefix sum greater than or equal to the random target. Using binary search on the prefix sums array reduces the pickIndex call time to O(log n)…

Full step-by-step walkthrough on Pro →

Optimal Approach

Preview

Precompute prefix sums of the weights during initialization. For each pickIndex call, generate a random integer target between 1 and total weight inclusive…

Full step-by-step walkthrough on Pro →

Want the full reasoning chain?

Unlock the complete walkthrough, line-by-line analysis, and recall drill.

Unlock Pro

Time

O(n) for initialization, O(log n) per pickIndex call

Computing prefix sums requires a single pass through the weights (O(n)). Each pickIndex call uses binary search on the prefix sums array, which takes O(log n) time.

Space

O(n)

The prefix sums array stores cumulative sums for all weights, requiring O(n) auxiliary space proportional to the input size.

Solution

Python

1class Solution:
2    def __init__(self, w: list[int]):
3        self.prefix = []
4        total = 0
5
6        for weight in w:
7            total += weight
8            self.prefix.append(total)
9
10        self.total = total
11
12    def pickIndex(self) -> int:
13        target = random.randint(1, self.total)
14
15        left = 0
16        right = len(self.prefix) - 1
17
18        while left < right:
19            mid = (left + right) // 2
20
21            if self.prefix[mid] < target:
22                left = mid + 1
23            else:
24                right = mid
25
26        return left

Step-by-Step Solution

1

Build Prefix Sum Array to Represent Weight Intervals

3self.prefix = []
4total = 0
6for weight in w:
7    total += weight
8    self.prefix.append(total)
10self.total = total

Objective

To create a prefix sums array that maps each index to the cumulative weight up to that index.

Key Insight

Prefix sums convert the discrete weights into contiguous intervals on the number line from 1 to total weight. This transformation allows the problem of weighted random selection to become a search for the interval containing a random target number. The prefix sums array is the foundation for efficient binary search.

Interview Quick-Check

Core Logic

Prefix sums represent cumulative weights, enabling mapping from a random target to an index by interval containment.

Common Pitfalls & Bugs

Forgetting to accumulate weights correctly or misaligning prefix sums with indices can cause incorrect interval mapping.

2

Generate Random Target and Use Binary Search to Find Corresponding Index

To pick a random integer in the total weight range and locate its corresponding index using binary search on prefix sums.

1 more step with full analysis available on Pro.

Line Analysis

This solution has 5 Critical lines interviewers watch for.

Line 26 Critical

return left

Return the left pointer as the index corresponding to the target.

After binary search completes, left points to the smallest prefix sum >= target, which maps to the weighted random index to return.

Line 8 Critical

self.prefix.append(total)

Append the current total to the prefix sums list.

Storing cumulative sums creates intervals that represent the weighted ranges for each index, which are essential for binary search.

Line 13 Critical

target = random.randint(1, self.total)

Generate a random integer target between 1 and total weight inclusive.

Randomly selecting a target in this range simulates picking an index with probability proportional to its weight by mapping the target to prefix sum intervals.

Full line-by-line criticality + rationale for all 16 lines available on Pro.

Test Your Understanding

Why is binary search used on the prefix sums array instead of a linear scan?

See the answer with Pro.

Don't just read it. Drill it.

Reconstruct Random Pick with Weight from memory until it sticks. AlgoDrill blanks out key lines and makes you fill them back in, step by step.

Unlock the Random Pick with Weight drill

or drill a free problem

1	`class Solution:`
2	`def __init__(self, w: list[int]):`
3	`self.prefix = []`
4	`total = 0`
5
6	`for weight in w:`
7	`total += weight`
8	`self.prefix.append(total)`
9
10	`self.total = total`
11
12	`def pickIndex(self) -> int:`
13	`target = random.randint(1, self.total)`
14
15	`left = 0`
16	`right = len(self.prefix) - 1`
17
18	`while left < right:`
19	`mid = (left + right) // 2`
20
21	`if self.prefix[mid] < target:`
22	`left = mid + 1`
23	`else:`
24	`right = mid`
25
26	`return left`

3	`self.prefix = []`
4	`total = 0`
6	`for weight in w:`
7	`total += weight`
8	`self.prefix.append(total)`
10	`self.total = total`