Problem Description

You need to design a data structure that supports weighted random selection from an array of positive integers.

Given an array w where w[i] represents the weight of index i, you need to implement a class with two functions:

1. Constructor __init__(self, w): Initializes the object with the array of weights.

2. Method pickIndex(): Returns a randomly selected index from the range [0, w.length - 1]. The probability of selecting each index i should be proportional to its weight, specifically w[i] / sum(w).

For example, if w = [1, 3]:

• Index 0 has weight 1, so its selection probability is 1/(1+3) = 0.25 (25%)
• Index 1 has weight 3, so its selection probability is 3/(1+3) = 0.75 (75%)

The solution uses a prefix sum array and binary search approach:

1. Initialization: Build a prefix sum array where s[i] contains the cumulative sum of weights from index 0 to i-1. This creates ranges for each index.

2. Random Selection:

   • Generate a random number x between 1 and the total sum of weights
   • Use binary search to find which range (and thus which index) this random number falls into
   • The index whose range contains x is returned

For instance, with w = [1, 3, 2]:

• Prefix sum array becomes [0, 1, 4, 6]
• Random numbers 1 map to index 0
• Random numbers 2-4 map to index 1
• Random numbers 5-6 map to index 2

This ensures each index is selected with probability proportional to its weight.

Intuition

The key insight is to transform the weighted selection problem into a simpler problem of selecting from continuous ranges.

Think of it like this: imagine you have a line segment of total length equal to the sum of all weights. Each index gets a portion of this line proportional to its weight. If we randomly throw a dart at this line, the index whose segment gets hit should be returned.

For example, with weights [1, 3, 2]:

• Total length = 6
• Index 0 occupies positions 1-1 (length 1)
• Index 1 occupies positions 2-4 (length 3)
• Index 2 occupies positions 5-6 (length 2)

To implement this efficiently, we use prefix sums to mark the boundaries of each segment. The prefix sum array [0, 1, 4, 6] tells us that:

• Index 0's range ends at position 1
• Index 1's range ends at position 4
• Index 2's range ends at position 6

When we pick a random number between 1 and 6:

• If it's 1, we're in index 0's range
• If it's 2, 3, or 4, we're in index 1's range
• If it's 5 or 6, we're in index 2's range

Since the prefix sum array is sorted, we can use binary search to quickly find which range contains our random number. We search for the smallest prefix sum that is greater than or equal to our random number - this tells us which index owns that position.

This approach converts the weighted probability problem into:

1. A preprocessing step to build ranges (O(n) time)
2. A binary search to find the right range (O(log n) per query)

The elegance lies in how prefix sums naturally create proportional ranges that match the required probabilities.

Learn more about Math, Binary Search and Prefix Sum patterns.

Solution Approach

The implementation consists of two main parts: initialization and random index selection.

Initialization (__init__ method)

We build a prefix sum array to create cumulative weight boundaries:

def __init__(self, w: List[int]):
    self.s = [0]
    for c in w:
        self.s.append(self.s[-1] + c)

• Start with self.s = [0] as the base
• For each weight c in the input array w, append the cumulative sum self.s[-1] + c
• For weights [1, 3, 2], this produces [0, 1, 4, 6]

This prefix sum array defines ranges:

• Index 0: range (0, 1]
• Index 1: range (1, 4]
• Index 2: range (4, 6]

Random Selection (pickIndex method)

def pickIndex(self) -> int:
    x = random.randint(1, self.s[-1])
    left, right = 1, len(self.s) - 1
    while left < right:
        mid = (left + right) >> 1
        if self.s[mid] >= x:
            right = mid
        else:
            left = mid + 1
    return left - 1

Step 1: Generate a random number x between 1 and the total sum (inclusive)

• random.randint(1, self.s[-1]) picks a random position on our "line"

Step 2: Binary search to find the smallest prefix sum ≥ x

• Initialize search bounds: left = 1, right = len(self.s) - 1
• We start from index 1 (not 0) because self.s[0] = 0 is just a placeholder
• The binary search finds the leftmost position where self.s[mid] >= x

Step 3: The loop maintains the invariant:

• If self.s[mid] >= x: the answer could be at mid or to its left, so right = mid
• If self.s[mid] < x: the answer must be to the right, so left = mid + 1
• The bit shift (left + right) >> 1 is equivalent to integer division by 2

Step 4: Return the result

• When the loop ends, left points to the position in the prefix sum array
• Return left - 1 to convert from prefix sum index to original array index

Time and Space Complexity

• Initialization: O(n) time to build the prefix sum array, O(n) space to store it
• pickIndex: O(log n) time for binary search, O(1) additional space
• The trade-off is preprocessing time for faster queries

Example Walkthrough

Let's walk through a concrete example with weights w = [2, 5, 3].

Initialization Phase:

1. Start with self.s = [0]
2. Process weight 2: self.s = [0, 2] (0 + 2)
3. Process weight 5: self.s = [0, 2, 7] (2 + 5)
4. Process weight 3: self.s = [0, 2, 7, 10] (7 + 3)

The prefix sum array [0, 2, 7, 10] creates these ranges:

• Index 0: numbers 1-2 (range size = 2, matching weight 2)
• Index 1: numbers 3-7 (range size = 5, matching weight 5)
• Index 2: numbers 8-10 (range size = 3, matching weight 3)

pickIndex() Call #1: Let's say the random number generated is x = 4

Binary search to find the smallest prefix sum ≥ 4:

• Initial: left = 1, right = 3 (searching positions 1-3 in prefix array)
• Iteration 1: mid = 2, self.s[2] = 7, since 7 ≥ 4, set right = 2
• Iteration 2: left = 1, right = 2, mid = 1, self.s[1] = 2, since 2 < 4, set left = 2
• Loop ends: left = right = 2
• Return left - 1 = 1

Result: Index 1 is selected (correct, since 4 falls in the range 3-7)

pickIndex() Call #2: Random number is x = 9

Binary search:

• Initial: left = 1, right = 3
• Iteration 1: mid = 2, self.s[2] = 7, since 7 < 9, set left = 3
• Loop ends: left = right = 3
• Return left - 1 = 2

Result: Index 2 is selected (correct, since 9 falls in the range 8-10)

Probability Verification:

• Index 0: selected when x ∈ [1,2], probability = 2/10 = 20%
• Index 1: selected when x ∈ [3,7], probability = 5/10 = 50%
• Index 2: selected when x ∈ [8,10], probability = 3/10 = 30%

These probabilities match the weight ratios exactly!

Solution Implementation

Time and Space Complexity

Time Complexity:

• __init__ method: O(n) where n is the length of the input array w. The method iterates through all elements once to build the cumulative sum array.
• pickIndex method: O(log n) where n is the length of the input array w. The method uses binary search on the cumulative sum array which has length n + 1.

Space Complexity:

• O(n) where n is the length of the input array w. The cumulative sum array self.s stores n + 1 elements (including the initial 0).

Overall Analysis:

• The initialization has a one-time cost of O(n) time and uses O(n) additional space to store the prefix sums.
• Each call to pickIndex() performs a binary search in O(log n) time.
• The algorithm trades space for time efficiency - by precomputing and storing cumulative sums, it enables fast weighted random selection through binary search rather than linear scanning.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Off-by-One Error in Random Number Generation

Pitfall: Using random.randint(0, self.prefix_sums[-1]) instead of random.randint(1, self.prefix_sums[-1]).

When you include 0 in the random range, it creates an edge case where no index would be selected properly since all prefix sums are greater than or equal to their first non-zero value. This breaks the probability distribution.

Incorrect:

# This can generate 0, which doesn't map correctly to any index
target = random.randint(0, self.prefix_sums[-1])

Correct:

# Start from 1 to ensure proper mapping
target = random.randint(1, self.prefix_sums[-1])

2. Wrong Binary Search Bounds

Pitfall: Starting binary search from index 0 instead of index 1.

Since prefix_sums[0] = 0 is just a placeholder and doesn't correspond to any actual weight range, including it in the binary search can lead to incorrect results.

Incorrect:

left, right = 0, len(self.prefix_sums) - 1  # Includes the placeholder 0

Correct:

left, right = 1, len(self.prefix_sums) - 1  # Skip the placeholder

3. Incorrect Binary Search Implementation

Pitfall: Using left = mid instead of left = mid + 1 when self.prefix_sums[mid] < target.

This can cause an infinite loop when left and right differ by 1, as mid would equal left and the bounds wouldn't change.

Incorrect:

if self.prefix_sums[mid] >= target:
    right = mid
else:
    left = mid  # Can cause infinite loop!

Correct:

if self.prefix_sums[mid] >= target:
    right = mid
else:
    left = mid + 1  # Always move left pointer forward

4. Using bisect Incorrectly

Alternative Pitfall: If using Python's bisect module, using bisect_left instead of bisect_right.

Since we want to find the smallest index where the prefix sum is greater than or equal to our target, we need the right behavior.

Incorrect:

import bisect
index = bisect.bisect_left(self.prefix_sums, target)
return index - 1  # May return wrong index for boundary cases

Correct:

import bisect
# bisect_right finds insertion point after any existing entries of target
index = bisect.bisect_left(self.prefix_sums, target + 1) 
# Or use bisect_right with target directly
index = bisect.bisect_right(self.prefix_sums, target)
return index - 1

5. Handling Edge Cases

Pitfall: Not considering arrays with single elements or zero weights.

While the problem states weights are positive integers, defensive programming suggests validating inputs:

Enhanced initialization:

def __init__(self, w: List[int]):
    if not w or not all(weight > 0 for weight in w):
        raise ValueError("Weights must be positive integers")
  
    self.prefix_sums = [0]
    for weight in w:
        self.prefix_sums.append(self.prefix_sums[-1] + weight)