Problem 91 · Data Engineering & Streaming

Reservoir Sampling

MediumClassic DSA & systems

01Problem

Return a uniform random sample of k items from a stream of unknown length in a single pass using O(k) memory.

02Key insight

Keep the first k items; for item i (0-based) replace a random slot with probability k/(i+1). By induction every item ends in the sample with probability k/n. Inject the RNG for reproducibility.

03Solution

solution.pypython16 lines

import random
from typing import Iterable, TypeVar

T = TypeVar("T")


def reservoir_sample(stream: Iterable[T], k: int, rng: random.Random = random.Random()) -> list[T]:
    sample: list[T] = []
    for i, item in enumerate(stream):
        if i < k:
            sample.append(item)
        else:
            j = rng.randrange(i + 1)  # uniform in [0, i]
            if j < k:
                sample[j] = item
    return sample

Try it first.

TimeO(n)

SpaceO(k)

tests.py · 3 passing

tests.pypython11 lines

import random
from collections import Counter
assert sorted(reservoir_sample([1, 2], 5)) == [1, 2]
rng = random.Random(42)
counts = Counter()
trials = 20000
for _ in range(trials):
    counts.update(reservoir_sample(range(10), 3, rng))
expected = trials * 3 / 10
assert all(abs(counts[i] - expected) < 0.06 * expected for i in range(10))
assert len(reservoir_sample(iter(range(10**4)), 7, rng)) == 7

04Staff-level follow-up

How do you sample across 1000 partitions in parallel? Tag each item with a random key u and keep the k smallest per partition, then merge by taking the global k smallest (equivalent to a uniform sample); for weighted sampling draw u in (0,1), keep the k LARGEST u^(1/w) (Efraimidis-Spirakis A-Res). Algorithm L cuts RNG calls to O(k log(n/k)).