Math, bits and matrices / Math / 135 · Random pick and reservoir sampling

Random pick and reservoir sampling

Prefix sums + binary search for weighted pick; reservoir keeps each item with prob 1/i.

The invariant

After i items, each has probability 1/i (reservoir) or weight/total (weighted).

01The intuition

Start with what you would do by hand

You have four raffle tickets with weights 1, 3, 2 and 4, and you want to draw one so that the ticket with weight 4 is four times as likely as the ticket with weight 1. By hand you would write the first ticket on 1 slip of paper, the second on 3 slips, the third on 2 and the fourth on 4, put all 10 slips in a hat, and draw one. That is correct. It is also wasteful: if the weights were 1 and 999,999 you would be writing a million slips.

Build the structure, one piece at a time

Instead of slips, lay the weights end to end on a number line. Ticket 0 owns the range 1, ticket 1 owns 2 to 4, ticket 2 owns 5 to 6, ticket 3 owns 7 to 10. The right edges of those ranges are the running totals, called prefix sums:

ticketweightprefix (running total)owns the numbers
0111
1342, 3, 4
2265, 6
34107, 8, 9, 10

Now draw a single random integer r from 1 to 10 and find which ticket owns it: the first ticket whose prefix is at least r. Because the prefix column is sorted, that lookup is a binary search. Draw r = 7: the middle prefix 4 is below 7, so look right; prefix 6 is below 7, look right; prefix 10 is the first one at or above 7, so ticket 3 wins. Draw r = 4: prefix 4 is the first one at or above 4, ticket 1 wins.

The second tool in this family handles the opposite situation, where you do not know how many items there are because they arrive one at a time and cannot be stored. Reservoir sampling keeps one chosen item. When item number i arrives, replace the chosen item with it with probability 1/i. One variable and a coin flip per item.

Why it is correct

For the weighted pick, ticket k owns exactly wk of the integers 1 … W, where W is the total, and each integer is equally likely, so

P(ticket k) = wk / W

which is the definition of a weighted draw. The binary search is correct because prefixes are non-decreasing, so "first prefix ≥ r" is a boundary with everything false on the left and true on the right.

For the reservoir, the invariant is that after i items each of them is the chosen one with probability 1/i. It is true for i = 1. When item i + 1 arrives it is chosen with probability 1/(i + 1). Any earlier item was chosen with probability 1/i and survives the new coin flip with probability i/(i + 1), so its chance is (1/i) · (i/(i + 1)) = 1/(i + 1), the same as the newcomer's. The invariant holds for every i, including the final count.

Why it is fast

Building the prefix table is one pass, O(n), done once. Each pick is a binary search over n prefixes, O(log n), and uses O(1) random numbers. The slips-in-a-hat method needed O(W) memory, and W can be astronomically larger than n. Reservoir sampling is O(1) memory and O(1) work per item, which is the best possible when the items cannot be stored.

How to recognise it in an interview

Cues: "pick an index with probability proportional to", "random node of a linked list in one pass", "shuffle", "random point inside these rectangles", or "the stream is too large to store". Say: "weighted means prefix sums plus a binary search; unknown length means reservoir sampling".

Random Pick with Weight is the table above exactly. Linked List Random Node is the reservoir: walk the list once, keeping each node with probability 1/i. Random Point in Non-overlapping Rectangles is the weighted pick where each rectangle's weight is its number of integer points, followed by a uniform point inside the chosen rectangle.

02The idea in one line

Prefix sums + binary search for weighted pick; reservoir keeps each item with prob 1/i.

Reach for it when: Random pick with weight, reservoir sampling on lists, shuffle, random pick index.

Cost: O(log n) per weighted pick.

03The walk-through, step by step

The video works one example: Random Pick with Weight w=[1,3]. These are its captions in order, so you can follow without the video or check a step you missed.

  1. Heavier items deserve a bigger slice of the random range
  2. prefix[0] = 0 + 1 = 1
  3. prefix[1] = 1 + 3 = 4
  4. prefix[2] = 4 + 2 = 6
  5. prefix[3] = 6 + 4 = 10
  6. Draw r = 7 uniformly from 1..10
  7. mid prefix 4 < 7: range becomes [2,3]
  8. mid prefix 6 < 7: range becomes [3,3]
  9. r = 7 lands on index 3 (weight 4, chance 4/10)
  10. Draw r = 1 uniformly from 1..10
  11. mid prefix 4 ≥ 1: range becomes [0,1]
  12. mid prefix 1 ≥ 1: range becomes [0,0]
  13. r = 1 lands on index 0 (weight 1, chance 1/10)
  14. Draw r = 4 uniformly from 1..10
  15. mid prefix 4 ≥ 4: range becomes [0,1]
  16. mid prefix 1 < 4: range becomes [1,1]
  17. r = 4 lands on index 1 (weight 3, chance 3/10)
  18. Prefix once O(n); each pick is a binary search O(log n)

04Asked at Google

7 Google-tagged LeetCode problems use this pattern. Frequency is LeetCode's Google-tag score (higher means reported more often); "last 6 months" marks problems reported in the most recent window of the export.

ProblemDifficultyFrequencyReported
Insert Delete GetRandom O(1)Medium46last 6 months
Random Pick with WeightMedium44last 6 months
Linked List Random NodeMedium29
Random Point in Non-overlapping RectanglesMedium21
Shuffle an ArrayMedium17last 6 months
Implement Rand10() Using Rand7()Medium15
Random Pick IndexMedium5

Other patterns in Math: Fast exponentiation, GCD and number theory, Digit-by-digit arithmetic, Coordinate geometry with hashing.