Reservoir Sampling: Uniform k-of-n Sample From a Stream of Unknown Length
Company: StackAdapt
Role: Machine Learning Engineer
Category: Statistics & Math
Difficulty: medium
Interview Round: Onsite
Items arrive one at a time from a stream. You do not know the total number of items, n, in advance, and the stream may be far too large to store. Design a one-pass algorithm that always holds a uniform random sample of k of the items seen so far, using memory proportional to k rather than n. Prove that when the stream ends, every item is in the sample with probability k/n.
Start with k = 1, then generalize to any k.
```hint Work backwards from the target
Right after the i-th item arrives, it must be in the sample with the same probability as every earlier item. Work out what that probability is, then ask what it forces the replacement rule to be.
```
```hint Survival is a product
For an item that is already in the sample, write down its chance of surviving one later arrival, then multiply those chances across all later arrivals.
```
### Constraints and Clarifications
- Only one pass over the stream is allowed, and items cannot be revisited.
- Memory is O(k) beyond the current item.
- A source of independent uniform random numbers is available.
- The sample is drawn without replacement. If the stream has fewer than k items, the sample is the whole stream.
### Clarifying Questions
- Must the sample be valid at every point in the stream, or only when the stream ends?
- Is sampling without replacement required, or would sampling with replacement do?
- Do items carry weights, or should every item be equally likely?
### What a Strong Answer Covers
- The exact replacement rule, including which slot is replaced
- A proof that each item's final inclusion probability is k/n, and ideally that every k-subset is equally likely
- Time per item, memory, and the number of random draws
- Edge cases: n < k, k = 1, and off-by-one errors in the random index range
### Follow-up Questions
- Items carry positive weights and should be included with probability that grows with their weight. How do you adapt the algorithm?
- The stream is split across several machines, each keeping its own reservoir. How do you merge them into one uniform sample of the whole stream?
- When n is much larger than k, most items are never kept. How can you avoid drawing a random number for every item?
Overview: Design a one-pass algorithm that keeps a uniform random sample of k items from a stream whose length is unknown in advance, using memory proportional to k. Tests the replacement rule, an inductive proof that every item is kept with equal probability, complexity, and edge cases.