Start from brute force
The obvious way to sample k of n items is to buffer the whole stream, then pick k at random at the end, for example with random.sample(items, k). That is exactly uniform and takes O(n) time, but it needs O(n) space and the end of the stream before it can answer. A firehose of logs does not fit in memory and may never end, so there is no moment at which to sample. The other instinct, keeping each item with probability k/n, needs n up front and returns a sample whose size varies from run to run instead of exactly k.
The approach
Decide each item’s fate the moment it arrives. The reference reservoir starts with res = stream[:k], which also covers a stream shorter than k. For every later index i it reads the supplied draw j = randoms[i - k], uniform over [0, i], and if j < k it writes res[j] = stream[i]. That admits item i with probability k/(i + 1), and since j is uniform over the k slots whenever it is below k, the evicted item is a uniformly chosen resident.
The invariant, by induction: after index i, each of the first i + 1 items is in the reservoir with probability k/(i + 1). It holds at i = k - 1, where everything is kept. At the next index a resident is evicted only if the new item is admitted and lands on its slot, probability (k/(i + 2)) · (1/k), so it stays with probability (i + 1)/(i + 2), and k/(i + 1) · (i + 1)/(i + 2) = k/(i + 2), the new item’s own chance. The TypeScript version runs the same loop over stream.slice(0, k).
Walk through an example
On the first example, stream = [1, 2, 3, 4, 5], k = 2, randoms = [0, 3, 1], two of the three later items get in.
| step | i | stream[i] | draw j | admit chance | result | reservoir |
|---|
| 1 | start | — | — | — | copy stream[:2] | [1, 2] |
| 2 | 2 | 3 | 0, from [0, 2] | 2/3 | 0 < 2: replace slot 0 | [3, 2] |
| 3 | 3 | 4 | 3, from [0, 3] | 2/4 | 3 is not below 2: skip | [3, 2] |
| 4 | 4 | 5 | 1, from [0, 4] | 2/5 | 1 < 2: replace slot 1 | [3, 5] |
| 5 | end | — | — | — | return | [3, 5] |
Complexity
Time O(n): copying the first k items is O(k), and every later item costs one comparison and at most one write. Space O(k) for the reservoir, whatever the stream length; the algorithm needs only the current item, although this problem passes the stream as a list. Buffering and then sampling is also O(n) time but takes O(n) space and must wait for the stream to end.
Edge cases to test
reservoir([1, 2], 5, []) → [1, 2]. A stream shorter than k is returned whole; stream[:k] handles it and no draw is read.reservoir([9, 8, 7], 3, []) → [9, 8, 7]. Exactly k items, so the loop never runs.reservoir([0, 1, 2, 3, 4], 2, [0, 0, 0]) → [4, 1]. The same slot can be overwritten again and again; slot 1 keeps its original item.reservoir([1, 2, 3], 2, [2]) → [1, 2], while the draw 1 gives [1, 3]. The test is strict: with j <= k, Python raises IndexError on res[2] and JavaScript silently grows the array to three items.reservoir(list(range(10)), 3, [3, 4, 5, 6, 7, 8, 9]) → [0, 1, 2]. Every draw is at least k, so nothing is replaced.reservoir([1, 2, 3], 0, [0, 1, 2]) → []. With k = 0 no draw is below k, so the reservoir stays empty.
Common mistakes
- Requiring the total count first — that defeats the entire point of the algorithm.
- Using the wrong replacement probability — for 0-based index
i it must be exactly k/(i + 1), which is what a draw from [0, i] landing below k gives. - Biasing toward early or late items by mishandling the boundary at
i = k.
Follow-ups an interviewer asks next
- Items have weights. Sample with probability proportional to weight.
- Use the Efraimidis–Spirakis scheme: give each item the key
u ** (1 / w), where u is uniform in (0, 1) and w is the item’s weight, and keep the k items with the largest keys in a min-heap of size k. Each item costs O(log k), so the pass is O(n log k) with O(k) memory. - The stream is enormous. Can you avoid one random draw per item?
- Yes. Algorithm L draws how many items to skip before the next replacement, from a geometric-style distribution, so it only touches the items that actually enter the reservoir. The expected number of random draws falls to O(k(1 + log(n/k))) instead of O(n), and the sample stays exactly uniform.
- Each shard of a cluster sees part of the stream. Build one global sample.
- Tag every item with an independent uniform random key and keep, on each shard, the k items with the smallest keys. The k smallest keys of the union are a uniform sample of everything, so merging is just taking the k smallest across the shard results. Plain per-shard reservoirs can also be merged, but only by weighting each shard by how many items it saw.
Where this shows up
Reservoir sampling draws a uniform random sample from a stream you cannot hold in memory or whose size you do not know — log sampling, A/B bucketing, picking a random line from a huge file, or sampling high-volume telemetry. It is the canonical “fair sample in one pass” technique.
Reference solution
Reference solution — Python
def reservoir(stream, k, randoms):
res = stream[:k] # start with the first k items
for i in range(k, len(stream)):
j = randoms[i - k] # provided draw in [0, i]
if j < k: # with prob k/(i+1): keep the new item
res[j] = stream[i]
return res
print(reservoir([1, 2, 3, 4, 5], 2, [0, 5, 1])) # -> [3, 5]
Reference solution — TypeScript
function reservoir(stream: number[], k: number, randoms: number[]): number[] {
const res = stream.slice(0, k); // first k items
for (let i = k; i < stream.length; i++) {
const j = randoms[i - k]; // provided draw in [0, i]
if (j < k) res[j] = stream[i]; // with prob k/(i+1): keep new item
}
return res;
}
console.log(reservoir([1, 2, 3, 4, 5], 2, [0, 5, 1])); // -> [3, 5]