Coding interview: Implement weighted reservoir sampling

1. Prompt and context

Each record in a stream has a positive weight, but neither the item count nor total weight is known. Implement a size-k weighted sample without replacement: each record appears at most once, its chance of inclusion is proportional to its weight, and the stream is scanned once. Explain key generation, candidate maintenance, extreme weights, and distribution tests.

2. What the interviewer is testing

• Whether you distinguish sampling with and without replacement and understand probability proportional to weight.
• Whether you can convert weighted sampling into keeping the top k random priorities or exponential keys.
• Whether you choose a size-k min-heap and give an O(log k) update per record.
• Whether you handle zero, huge, or tiny weights, random boundaries, duplicate IDs, and a reproducible seed.

3. Clarifying questions before answering

1. Are weights finite positive numbers, and should zero-weight records be dropped or retained?
2. Is the sample without replacement, or may a record appear more than once?
3. Is only the final sample required, or must every prefix have the correct distribution?
4. Must multiple shards be merged, state persisted, or results reproduced exactly?

4. A 30-second answer framework

Generate an independent random key for each record of weight w and keep the largest k keys. A stable form draws u uniformly from (0, 1] and computes key = log(u) / w; because keys are negative, this is equivalent to keeping the k keys closest to zero. Store the current sample in a size-k min-heap whose root is the smallest key; replace it only when a new key is larger. A single pass takes O(n log k) time and O(k) extra space. Validate weights and make the random source injectable.

5. Step-by-step deep answer

Step 1: Define the distribution

Weighted sampling without replacement is not k independent draws with probability w / total, because that can repeat a record. The target is a size-k set whose order-statistic distribution equals repeatedly drawing an unselected item in proportion to its remaining weight. The reservoir should be a valid sample after every prefix, not only after the stream ends.

Step 2: Generate a numerically stable key

The exponential race gives a convenient implementation: draw a uniform u and calculate key = log(u) / w, then retain the largest keys. As u approaches zero, log(u) becomes more negative; a larger weight makes the key closer to zero and therefore more likely to enter the top k. Avoid u ** (1 / w), which can underflow or lose separation for extreme weights.

sample_key(weight):
    require finite(weight) and weight > 0
    u = uniform_random_open_interval()
    return log(u) / weight

Step 3: Maintain top-k with a min-heap

Store (key, sequence, item) in the heap, using sequence to break exact key ties. Push while the reservoir is not full. Once full, compare the new key with the root and replace only when it is larger. If k is zero, discard every record. Re-sorting an array after every record would make updates O(k log k) instead of O(log k).

Step 4: Handle input and randomness boundaries

Reject NaN, infinity, and negative weights. A zero-weight record cannot be selected by a positive-weight sample and can be skipped. The random source must not return zero, or log(0) becomes unusable; redraw or clamp to the smallest positive floating-point value. Treat duplicate IDs as separate records unless the problem explicitly asks for ID deduplication. Inject a pseudo-random source in tests so failures are reproducible.

Step 5: Complexity, validation, and distributed extension

For n records, the single-machine implementation takes O(n log k) time and O(k) extra space. Use fixed-weight Monte Carlo simulation to check that marginal inclusion rises with weight, and assert that the sample has no duplicates. In a distributed stream, each shard can generate keys with the same rule and a coordinator can merge the shard top-k candidates; state, updates, deletions, seeds, and communication cost still need an explicit design. Simply sampling uniformly from the shard reservoirs loses discarded-record information.

6. Example of a high-quality answer

For every positive-weight record I generate key = log(u) / w, where u is uniform on an open interval, and keep the largest k keys. Keys are negative, so larger weights tend to be closer to zero. A capacity-k min-heap stores the sample; when full, a new key replaces the smallest root only if it is larger. I define behavior for invalid weights, zero random values, k = 0, and duplicate records. The scan is O(n log k) time and O(k) space. I validate the distribution with repeated fixed-weight simulations and merge distributed candidates by the same global top-k key rule.

7. Common mistakes

• Drawing independently with w / total → creates duplicates and is not without-replacement sampling → use random keys and top-k.
• Computing u ** (1 / w) directly → underflow for extreme weights → compare logarithmic keys.
• Using a max-heap for top-k → requires searching for the smallest value → use a min-heap so the replacement point is the root.
• Allowing u = 0 → log(0) becomes negative infinity → use an open-interval source or redraw.
• Testing one output only → misses long-run bias → run fixed-weight Monte Carlo tests and assert no duplicates.

8. Follow-ups and responses

Why does key = log(u) / w give weighted sampling?

Treat -log(u) as an exponential variable with rate one. Dividing by w gives an exponential time with rate w. The smallest exponential time is more likely to come from a larger rate; negating it means retaining the largest keys, which yields weighted sampling without replacement.

How do you make results reproducible?

Inject a pseudo-random source with an explicit seed and record the item identifier, weight version, and algorithm version in experiment metadata. Do not depend on thread scheduling or global random state, or identical input may produce different samples.

Can two already-built reservoirs be merged?

If both shards generated independent keys under the same rule, merge their candidate keys and take the global top k. Treating only the two final reservoirs as ordinary data and sampling again loses information about discarded records. Persisted keys and shard updates or deletions also need defined behavior.

What if weights change over time?

Changing a weight changes the target distribution, so the old key no longer represents the new weight. Regenerate keys for affected records or re-enter a weight-versioned event into the stream. State whether a short approximation is acceptable and how old samples are retired.