Reservoir Sampling

Reservoir sampling picks a uniformly random item from a stream whose length you do not know in advance — and it does it in a single pass using only O(1) extra memory. The trick is to keep one "reservoir" slot and, as each new item arrives, decide on the spot whether it should replace what you are holding.

The classic setup: a log file too big to fit in memory, or an endless stream, and you want one fair random sample. For stream = [5, 2, 8, 1, 9] you keep item 1 with probability 1, swap to item 2 with probability 1/2, to item 3 with probability 1/3, and so on — and any of the five values is equally likely to remain.

Intuition

Why does 1/i give a uniform result? Look at the last item. It replaces whatever you held with probability 1/n, so it survives exactly 1/n of the time — fair. Now take any earlier item x at position i. It first has to win its own slot (probability 1/i), and then it has to survive every later item. Item i+1 leaves it alone with probability i/(i+1), item i+2 with probability (i+1)/(i+2), and so on. Multiply them:

(1/i) · (i/(i+1)) · ((i+1)/(i+2)) · … · ((n-1)/n) = 1/n.

Every numerator cancels the previous denominator — a telescoping product that collapses to 1/n for every position. That cancellation is the whole proof.

Walk through it

Step through the animation on the right. The i pointer scans the stream left to right; the single res cell below holds the current pick, and the label shows the replacement probability for each item.

When i = 1, the coin is 1/1 = 1, so res simply becomes 5. At i = 2 the 1/2 flip lands on "keep", so res stays 5. At i = 3 the 1/3 flip lands on "replace" — res lights up and takes 8. At i = 4 (1/4) and i = 5 (1/5) the flips both land on "keep", so 8 survives to the end. The final sample is 8, but remember: in a real run each of the five values would come out 1/5 of the time.

The code, line by line

import random
 
def reservoir_sample(stream):
    res = None
    for i, x in enumerate(stream, 1):
        if random.randint(1, i) == 1:
            res = x
    return res

• res = None is the reservoir slot — we hold exactly one item, never the whole stream.
• enumerate(stream, 1) walks the items while counting from 1, so i is the 1-indexed position we need for the 1/i probability.
• Line 6 is the coin flip: random.randint(1, i) is uniform over 1..i, so == 1 is true with probability exactly 1/i.
• Line 7 performs the replacement when the flip wins; otherwise we keep what we had.
• return res hands back the single survivor after one pass — no second loop, no buffering.

Complexity

The win is the space bound: you never store the stream, so this works on inputs far larger than memory, and even on streams whose length you discover only when they end. Generalizing to k > 1 keeps a reservoir of size k and replaces a random slot with probability k/i — same O(k) space, same single pass.

When to use / pitfalls

Practice