Reservoir sampling

Python data engineering interview problem. Difficulty: advanced. Pattern: Streaming. About 20 minutes. Part of the Pro drill bank.

Pick k items uniformly at random from a stream of unknown length, in one pass and O(k) memory. Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

A stream can be too large to store, yet a job still needs a fair sample of k items with repeatable tests. Read the iterable once and choose k items uniformly, using only O(k) memory and random.Random(seed). If the stream is shorter than k, return all items in stream order. Write reservoir_sample(stream, k: int, seed: int) -> list. The first sample asks whether item 0 appears fewer than 400 times in 1000 seeded samples; the correct algorithm returns true because inclusion is about 30 percent. The starter file already contains the function name, harness and print lines. Keep them and fill only the function body.

Requirements

  • Consume the stream only once.
  • Keep no more than k sampled items.
  • Use the supplied seed for repeatable random choices.

Constraints

  • The stream length is unknown and may be a generator.
  • Every item should have chance k / n when n is at least k.

Examples

Input: sum(0 in reservoir_sample(iter(range(10)), 3, seed) for seed in range(1000)) < 400 Output: True A fair sample of 3 out of 10 contains item 0 in about 3 of every 10 runs (roughly 300 of 1000), well under 400.

Topics: lakebench, python, sampling, stream, random.

More Python interview questions · All interview problems · Learn data engineering

advanced

Reservoir sampling

Interview-style drill: Pick k items uniformly at random from a stream of unknown length, in one pass and O(k) memory.

Implement `reservoir_sample(stream, k: int, seed: int) -> list`. `stream` is an iterable you can loop over **once**, and its length is not known in advance (it may be a generator). Return a list of `k` items chosen uniformly at random from it: every item must have the same chance, k / n, of being in the result. If the stream has fewer than `k` items, return all of them in stream order. Use `random.Random(seed)` for all random choices so the result is repeatable for the same seed. Do not store the whole stream.