528. Random Pick with Weight

528

MediumRandomization, Math and Miscellaneous (FAANG add-on)Weighted random sampling via cumulative distributionPrefix sums + binary search

Random Pick with Weight

Given an array w of positive integer weights, implement pickIndex() that returns index i with probability w[i] divided by the sum of all weights.

Open official problem prompt ↗

In plain English

Choose an index at random so that each index's chance is proportional to its weight, efficiently over many queries.

Picture it like this

A raffle where index i gets w[i] tickets in one big drum; you draw a single ticket number and see which index's block of tickets it falls into.

Example

Input: w = [1,3]; call pickIndex() repeatedly
Output: [1]
Why: Index 0 is returned with probability 1/4 and index 1 with probability 3/4; a single call here returned 1, its more likely outcome.

Constraints

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

Pattern lesson

See the pattern, then code

Prefix sums + binary search

↗

Randomization, Math and Miscellaneous (FAANG add-on)

Uniform random selection, sampling streams, number theory, or precision-heavy parsing.. Lay the weights end to end on a line of length total. A uniform random point in [1,total] lands in segment i with probability w[i]/total. Precompute prefix sums once; each pick is a uniform draw plus a binary search for the segment containing it.

New words, made simpleKnow these before the algorithm

Prefix sum: prefix[i] is the cumulative weight through index i, marking the right edge of segment i.
Cumulative distribution: The prefix array acts as a discrete CDF; binary search inverts it.
bisect_left: Returns the first position whose value is >= target, i.e. the segment that contains target.

Approaches from first idea to best ideaCompare the trade-offs

ApproachHow it thinksCost

Expand into a big list

Memory blows up: weights up to 10^5 across 10^4 entries is up to 10^9 cells.

Insert index i exactly w[i] times, then random.choice.

Time O(1) per pickSpace O(sum of weights)

Linear scan over prefix

Correct but slow when there are many picks and large n.

Draw target, walk the prefix array until prefix[i] >= target.

Time O(n) per pickSpace O(n)

Best choicePrefix sums + binary search

Compact and fast for up to 10^4 queries.

Precompute prefix sums, draw target, binary search for the segment.

Time O(log n) per pickSpace O(n)

The rule we keep true

Invariant

prefix is nondecreasing and prefix[i] - prefix[i-1] = w[i], so the half-open block of integers ending at prefix[i] has exactly w[i] members.

Why this is correct

Reasoning

target is uniform over the total integers 1..total. Exactly w[i] of those integers fall in index i's block (between prefix[i-1]+1 and prefix[i]), and bisect_left maps each such target to i, so P(i) = w[i]/total.

The algorithm in three movesSay these aloud before coding

1In the constructor build prefix sums where prefix[i] = w[0]+...+w[i]

prefix = [1,4], total=4

2Store the total sum

target = randint(1,4) = 3

3For pickIndex draw target = random integer in [1,total]

bisect_left([1,4],3) = 1 -> index 1

4Binary search (bisect_left) for the first prefix >= target and return that index

Visual trace

Follow the data, one decision at a time

Input snapshotHighlighted cells are involved in this example

1 · Readw=[1,3]

2 · AskCumulative weights?

3 · Update stateprefix=[1,4], total=4

4 · Resultsegments ready

Key takeaway

Prefix sums split [1,4] into segment [1,1] for index 0 and [2,4] for index 1; a random target lands proportionally.

Code walkthrough

Read the solution in small chunks

Python 3

Do not memorize the whole program. Connect each group of lines to one job in the algorithm.

1
Lines 6-11Prefix sums
Accumulate weights once so each query is cheap; store total for the draw range.
2
Lines 13-15Weighted draw
A uniform target in [1,total] plus bisect_left finds the containing segment in O(log n).

Make it stick

Edge cases, mistakes, and one self-check

Edge cases to test

Single weight — total equals that weight and index 0 is always returned
Very skewed weights (one huge, others tiny) still sampled correctly
Large total near 10^9 — Python ints handle it without overflow

Common beginner mistakes

Drawing target in [0,total] or [0,total-1] misaligns segments with bisect_left; use [1,total]
Using bisect_right instead of bisect_left double-counts boundaries and biases results
Rebuilding prefix sums on every pick instead of once in the constructor

Check your understanding

Why draw target in [1, total] and use bisect_left rather than [0, total-1]?

Random Pick with Weight

Random Pick with Weight

See the pattern, then code

Invariant

Reasoning

Follow the data, one decision at a time

Read the solution in small chunks

Edge cases, mistakes, and one self-check

Edge cases to test

Common beginner mistakes

More in Randomization, Math and Miscellaneous (FAANG add-on)