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 - LeetCode Solution

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

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). This approach balances preprocessing cost with efficient repeated queries.

Optimal Approach

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.

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

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

13target = random.randint(1, self.total)
15left = 0
16right = len(self.prefix) - 1
18while left < right:
19    mid = (left + right) // 2
21    if self.prefix[mid] < target:
22        left = mid + 1
23    else:
24        right = mid
26return left

Objective

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

Key Insight

Randomly selecting a number between 1 and total weight simulates weighted probability. Binary search efficiently finds the smallest prefix sum greater than or equal to this target, which corresponds to the weighted index. This approach avoids linear scanning and ensures logarithmic time complexity per pick.

Interview Quick-Check

Core Logic

Binary search finds the first prefix sum >= target, mapping the random target to the correct weighted index.

State & Boundaries

The search space is between left=0 and right=len(prefix)-1, and the loop continues while left < right to narrow down the interval.

Common Pitfalls & Bugs

Incorrect mid calculation or updating pointers improperly can cause infinite loops or off-by-one errors.

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.

Line 21 Critical

if self.prefix[mid] < target:

If the prefix sum at mid is less than the target, move left pointer rightward.

Since prefix sums are sorted, if prefix[mid] < target, the target must be in the right half, so left is updated to mid + 1.

Line 23 Critical

else:

Otherwise, move right pointer to mid.

If prefix[mid] >= target, the target interval could be at mid or to the left, so right pointer is moved to mid to narrow search.

Line 7 Core

total += weight

Add the current weight to the running total.

Updating the total weight at each step ensures prefix sums reflect cumulative weights up to the current index.

Line 10 Core

self.total = total

Store the total weight for use in random target generation.

The total weight defines the upper bound for the random integer, ensuring the random selection covers the entire weighted range.

Line 22 Core

left = mid + 1

Update left pointer to mid + 1 when prefix[mid] is less than target.

Advancing left pointer excludes the left half, focusing search on the right half where the target lies.

Line 24 Core

right = mid

Update right pointer to mid to narrow the search space.

This ensures the binary search converges on the smallest index with prefix sum >= target, which corresponds to the weighted index.

Line 3 Core

self.prefix = []

Initialize an empty list to store prefix sums.

This list will hold cumulative sums of weights, enabling efficient interval mapping for weighted random selection.

Line 4 Core

total = 0

Initialize a variable to accumulate total weight.

Tracking the total weight is necessary to define the range for random target generation and to build prefix sums correctly.

Line 18 Core

while left < right:

Continue binary search while left pointer is less than right pointer.

This loop condition ensures the search space is reduced until the exact interval is found, preventing infinite loops.

Line 19 Core

mid = (left + right) // 2

Calculate the middle index of the current search space.

Midpoint calculation divides the search space to efficiently narrow down the target interval.

Line 15 Standard

left = 0

Initialize left pointer for binary search at the start of prefix sums.

Left pointer marks the lower bound of the search space for locating the target interval.

Line 16 Standard

right = len(self.prefix) - 1

Initialize right pointer for binary search at the end of prefix sums.

Right pointer marks the upper bound of the search space, enabling narrowing down the interval containing the target.

Line 6 Standard

for weight in w:

Iterate over each weight in the input list.

Processing each weight sequentially allows building the prefix sums array cumulatively.

Test Your Understanding

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

Binary search reduces the time complexity of locating the target interval from O(n) to O(log n), which is critical for efficient repeated calls to pickIndex, especially when the input size is large.

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`