[Medium] 528. Random Pick with Weight

You are given a 0-indexed array of positive integers w where w[i] describes the weight of the ith index.

You need to implement the function pickIndex(), which randomly picks an index in the range [0, w.length - 1] (inclusive) and returns it. The probability of picking an index i is w[i] / sum(w).

For example, if w = [1, 3], the probability of picking index 0 is 1 / (1 + 3) = 0.25 (i.e., 25%), and the probability of picking index 1 is 3 / (1 + 3) = 0.75 (i.e., 75%).

Examples

Example 1:

Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1]],[],[],[],[],[]]
Output
[null,0,0,0,0,0]

Explanation
Solution solution = new Solution([1]);
solution.pickIndex(); // return 0. The only option is to return 0 since there is only one element in w.

Example 2:

Input
["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
[[[1,3]],[],[],[],[],[]]
Output
[null,1,1,1,1,0]

Explanation
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // return 1. It's returning the right index (1), and the probability of returning 1 is 3/4 = 0.75.
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 1
solution.pickIndex(); // return 0. The probability of returning 0 is 1/4 = 0.25.

Constraints

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

Thinking Process

Prefix Sum Creates Ranges: Each index gets a range proportional to its weight

The search space must shrink monotonically each step.
Decide which half still satisfies the predicate, discard the other.
Use mid = left + (right - left) / 2 to avoid overflow.

Common Approaches

Typical techniques for this pattern:

Approach	Time	Space	Notes
Standard binary search (this problem)	O(log n)	O(1)	Sorted array, `left <= right`
Lower / upper bound	O(log n)	O(1)	First/last position, insert index
Binary search on rotated array	O(log n)	O(1)	Identify sorted half, discard other
Binary search on answer	O(n log M)	O(1)	Monotonic predicate over search space

Solution

Time Complexity:

Constructor: O(n) - Build prefix sum array
pickIndex: O(n) linear search or O(log n) binary search

Space Complexity: O(n) - Prefix sum array

The key insight is to use prefix sums to create ranges proportional to weights, then use binary search to find the index corresponding to a random value.

Solution 1: Linear Search

import random


class Solution:
    def __init__(self, w: list[int]):
        self.prefix: list[int] = []
        s = 0
        for x in w:
            s += x
            self.prefix.append(s)

    def pickIndex(self) -> int:
        total = self.prefix[-1]
        target = random.random() * total
        for i, p in enumerate(self.prefix):
            if target < p:
                return i
        return len(self.prefix) - 1

Solution Explanation

Approach: Standard binary search (this problem)

Key idea: 1. Prefix Sum Creates Ranges: Each index gets a range proportional to its weight

How the code works:

Prefix Sum Creates Ranges: Each index gets a range proportional to its weight
- The search space must shrink monotonically each step.
- Decide which half still satisfies the predicate, discard the other.
- Use mid = left + (right - left) / 2 to avoid overflow.

Approach	Constructor	pickIndex	Space	Pros	Cons
Linear Search	O(n)	O(n)	O(n)	Simple	Slow for many calls
Binary Search (lower_bound)	O(n)	O(log n)	O(n)	Fast, clean	Requires STL
Custom Binary Search	O(n)	O(log n)	O(n)	Fast, explicit	More code

Solution 2: Binary Search (Optimized)

import bisect
import random


class Solution:
    def __init__(self, w: list[int]):
        self.prefix: list[int] = []
        s = 0
        for x in w:
            s += x
            self.prefix.append(s)

    def pickIndex(self) -> int:
        total = self.prefix[-1]
        target = random.random() * total
        return bisect.bisect_left(self.prefix, target)

Step-by-Step Example: `w = [1, 3, 2]`

Constructor:
  prefixSum = []
  n=1: prefixSum.push_back(1 + 0) = [1]
  n=3: prefixSum.push_back(3 + 1) = [1, 4]
  n=2: prefixSum.push_back(2 + 4) = [1, 4, 6]
  
  Ranges:
    Index 0: [0, 1)     → weight 1
    Index 1: [1, 4)     → weight 3
    Index 2: [4, 6)     → weight 2
    Total: 6

pickIndex() call 1:
  randNum = 0.7 (example)
  target = 0.7 * 6 = 4.2
  
  Linear Search:
    i=0: 4.2 < 1? No
    i=1: 4.2 < 4? No
    i=2: 4.2 < 6? Yes → return 2
    
  Binary Search:
    lower_bound([1,4,6], 4.2) → points to index 2
    return 2

pickIndex() call 2:
  randNum = 0.3 (example)
  target = 0.3 * 6 = 1.8
  
  Linear Search:
    i=0: 1.8 < 1? No
    i=1: 1.8 < 4? Yes → return 1
    
  Binary Search:
    lower_bound([1,4,6], 1.8) → points to index 1
    return 1

Visual Representation

Weights:     [1,  3,  2]
Prefix Sums: [1,  4,  6]

Range Mapping:
  0    1    4    6
  |----|----|----|
   [0]  [1]  [2]
   
Random value 0.0-6.0 maps to:
  [0.0, 1.0) → Index 0 (weight 1)
  [1.0, 4.0) → Index 1 (weight 3)
  [4.0, 6.0) → Index 2 (weight 2)

Algorithm Breakdown

Constructor

def __init__(self, w: list[int]):
    self.prefix: list[int] = []
    s = 0
    for x in w:
        s += x
        self.prefix.append(s)

How it works:

Build cumulative prefix sum array
prefixSum[i] = sum of weights from index 0 to i
Example: w = [1,3,2] → prefixSum = [1,4,6]

pickIndex - Linear Search

import random


def pickIndex(self) -> int:
    total = self.prefix[-1]
    target = random.random() * total
    for i, p in enumerate(self.prefix):
        if target < p:
            return i
    return len(self.prefix) - 1

How it works:

Generate random float in [0.0, 1.0)
Scale to [0.0, totalSum)
Find first prefix sum >= target
Return corresponding index

pickIndex - Binary Search

import bisect
import random


def pickIndex(self) -> int:
    total = self.prefix[-1]
    target = random.random() * total
    return bisect.bisect_left(self.prefix, target)

How it works:

lower_bound finds first element >= target
Returns iterator, subtract begin() to get index
O(log n) instead of O(n)

Complexity

| Approach | Constructor | pickIndex | Space | Pros | Cons | |———-|————-|———–|——-|——|——| | Linear Search | O(n) | O(n) | O(n) | Simple | Slow for many calls | | Binary Search (lower_bound) | O(n) | O(log n) | O(n) | Fast, clean | Requires STL | | Custom Binary Search | O(n) | O(log n) | O(n) | Fast, explicit | More code |

Implementation Details

Prefix Sum Construction

import bisect
import random


class Solution:
    def __init__(self, w: list[int]):
        self.prefix: list[int] = []
        s = 0
        for x in w:
            s += x
            self.prefix.append(s)

    def pickIndex(self) -> int:
        total = self.prefix[-1]
        target = random.randint(1, total)
        return bisect.bisect_left(self.prefix, target)

Why this works:

First element: n + 0 = n
Subsequent: n + previous_sum = cumulative_sum
Creates ranges: [0, prefixSum[0]), [prefixSum[0], prefixSum[1]), ...

Random Number Generation

import random


class Solution:
    def __init__(self, w: list[int]):
        self.prefix: list[int] = []
        s = 0
        for x in w:
            s += x
            self.prefix.append(s)

    def pickIndex(self) -> int:
        total = self.prefix[-1]
        target = random.randint(1, total)
        left, right = 0, len(self.prefix) - 1
        while left < right:
            mid = left + (right - left) // 2
            if self.prefix[mid] < target:
                left = mid + 1
            else:
                right = mid
        return left

Why float?

rand() / RAND_MAX gives uniform distribution in [0, 1)
Multiplying by total sum scales to [0, totalSum)
Float provides fine-grained selection

lower_bound Usage

s = 0
self.prefix = []
for n in w:
    s += n
    self.prefix.append(s)

What it does:

Returns iterator to first element >= target
If all elements < target, returns end()
Subtract begin() to get index

Common Mistakes

Single weight: w = [1] → always return 0
Equal weights: w = [1,1,1] → equal probability for all
Very large weight: w = [1, 1000000] → index 1 almost always picked
Single large weight: w = [100] → always return 0
Wrong random range: Using rand() % totalSum instead of scaling properly
Off-by-one errors: Not handling edge cases correctly
Integer overflow: Not considering large weights
Wrong comparison: Using <= instead of < in linear search
Not seeding random: Should call srand(time(nullptr)) in constructor

Optimization Tips

Use Binary Search: Always prefer binary search for O(log n) performance
Integer Random: Use rand() % totalSum + 1 for integer precision
Precompute Total: Store totalSum instead of calling prefixSum.back() repeatedly
Modern Random: Consider <random> header for better randomness

497. Random Point in Non-overlapping Rectangles - Similar weighted selection
710. Random Pick with Blacklist - Random with exclusions
380. Insert Delete GetRandom O(1) - Random from set
398. Random Pick Index - Random index of target value

Real-World Applications

Load Balancing: Weighted random selection of servers
Game Development: Weighted loot drops, random encounters
Sampling: Proportional sampling from populations
Recommendation Systems: Weighted content selection
A/B Testing: Weighted traffic distribution

Pattern Recognition

This problem demonstrates the “Weighted Random Selection” pattern:

Build prefix sum array (creates proportional ranges)
Generate random number in [0, totalSum)
Binary search to find corresponding index
Return index

Why Prefix Sum Works

Proportional Ranges: Each index gets range size = its weight
Non-overlapping: Ranges don’t overlap, covering [0, totalSum)
Uniform Random: Random number uniformly distributed maps to proportional selection
Efficient Lookup: Binary search finds index in O(log n)

Mathematical Proof

For weights w = [w₀, w₁, ..., wₙ₋₁]:

Total sum: S = Σwᵢ
Prefix sums: P = [w₀, w₀+w₁, ..., S]
Random value r ∈ [0, S) maps to index i where P[i-1] ≤ r < P[i]
Probability of r falling in range [P[i-1], P[i]) = wᵢ/S ✓

Random Number Generation Notes

Using `rand()`

import random

total = self.prefix[-1]
target = random.random() * total  # uniform in [0, total)

Pros:

Simple, built-in
Fast

Cons:

Lower quality randomness
Not thread-safe

Using `<random>` (Modern C++)

import bisect

idx = bisect.bisect_left(self.prefix, target)

Pros:

Better randomness quality
Thread-safe
More control

Cons:

Requires Python11+
More complex

This problem is an excellent example of combining prefix sums with binary search to achieve efficient weighted random selection, a common pattern in system design and algorithms.

Key Takeaways

Prefix Sum Creates Ranges: Each index gets a range proportional to its weight
Random Target: Generate random number in [0, totalSum) range
Binary Search: Find first prefix sum >= target (O(log n))
Linear Search: Simpler but slower (O(n))
Probability Proportional: Larger weights get larger ranges

References

LC 528: Random Pick with Weight on LeetCode
LeetCode Discuss — LC 528: Random Pick with Weight
LeetCode Editorial (may require premium)

Template Reference

Data Structure Design

Robina Li

[Medium] 528. Random Pick with Weight

Examples

Constraints

Thinking Process

Common Approaches

Solution

Solution 1: Linear Search

Solution Explanation

Solution 2: Binary Search (Optimized)

Step-by-Step Example: `w = [1, 3, 2]`

Visual Representation

Algorithm Breakdown

Constructor

pickIndex - Linear Search

pickIndex - Binary Search

Complexity

Implementation Details

Prefix Sum Construction

Random Number Generation

lower_bound Usage

Common Mistakes

Optimization Tips

Real-World Applications

Pattern Recognition

Why Prefix Sum Works

Mathematical Proof

Random Number Generation Notes

Using `rand()`

Using `<random>` (Modern C++)

Key Takeaways

References

Template Reference

Related Posts

Recent Posts

Examples

Constraints

Thinking Process

Common Approaches

Solution

Solution 1: Linear Search

Solution Explanation

Solution 2: Binary Search (Optimized)

Step-by-Step Example: w = [1, 3, 2]

Visual Representation

Algorithm Breakdown

Constructor

pickIndex - Linear Search

pickIndex - Binary Search

Complexity

Implementation Details

Prefix Sum Construction

Random Number Generation

lower_bound Usage

Common Mistakes

Optimization Tips

Related Problems

Real-World Applications

Pattern Recognition

Why Prefix Sum Works

Mathematical Proof

Random Number Generation Notes

Using rand()

Using <random> (Modern C++)

Key Takeaways

References

Template Reference

Related Posts

Recent Posts

Step-by-Step Example: `w = [1, 3, 2]`

Using `rand()`

Using `<random>` (Modern C++)