~/problems / Probability / Randomized algorithms

Uniform sample of k items from a stream

medium ~20 min

Write sample_stream(stream, k: int, rng: random.Random) -> list.

stream is an iterator of unknown length that you can read only once, and it may be far too long to store. Return a list of k items chosen uniformly at random without replacement: every set of k positions must be equally likely, so each item ends up in the result with probability k / n, where n is the stream's length. If the stream has fewer than k items, return all of them. The order of the returned list doesn't matter.

Use rng (e.g. rng.randrange(m)) for every random choice, never the global random module, so the tests can seed it.

rng = random.Random(0)
sample_stream(iter([10, 20]), 5, rng)          # [10, 20]  (fewer than k items)
sample_stream((x * x for x in range(100)), 3, rng)   # e.g. [49, 1024, 4]

Constraints: 1 <= k <= 1000; the stream may hold millions of items. Use O(k) memory: the tests measure it, so copying the stream into a list first fails.

Show hint

Keep a list of k items as your current sample. Decide what to do with item number i so that, after reading it, each of the first i items is in the sample with probability k / i; then check that this invariant survives each later step.

Topic: Randomized algorithms. Reservoir sampling, Fisher-Yates, rejection sampling.

0:00
Ctrl ' run · Ctrl ↵ submit
esc