Stream Sampling & Distribution Verification

Implement reservoir sampling over an unknown-size stream under O(k) memory and one-pass constraints, then write a simulation that verifies the selected items are close to uniformly distributed. The follow-up focuses on proving correctness empirically with repeated trials and distribution checks. Problem Requirements You have a stream of numbers (Data Stream) that is very large or infinite. You do not know how many numbers there are in total. You cannot fit all the numbers into your computer's memory. You need to pick k items from this stream randomly (for example, k=3). Rules: Memory Limit: You cannot save all the numbers. You can only store a small amount, roughly O(k). Single Pass: You can only look at the numbers one time. You cannot go back. Uniform Distribution: Every number in the stream must have the same chance of being picked. If there are N total numbers, the chance should be k/N. This interview question usually has two parts: Coding: Write the algorithm. Testing: Prove the algorithm works using code. Part 1: The Solution Strategy: Reservoir Sampling Because we do not know the total size N, we cannot use a standard random choice function. We must use an algorithm called Reservoir Sampling. How it works: Start: Create a list (the "reservoir") to hold k items. Fill the Start: Put the first k numbers from the stream directly into your list. Process the Rest: For every new number (let's call it the n-th number) that comes after: Pick a random number j between 1 and n. If j is less than or equal to k (this happens with probability k/n): Swap the n-th number with the element currently at index j-1 in your list. Otherwise, ignore the n-th number. Why does this give every number an equal chance? We can prove this step-by-step (Induction): Goal: We want to prove that after n numbers, the chance of any number being in our list is k/n. Step 1 (Base Case): When n = k, we keep all k items. The chance is k/k = 1. This is correct. Step 2 (The Next Number): Suppose we have n-1 numbers, and the chance for each was k/(n-1). Now we look at the new n-th number. The new number: We keep it with probability k/n. The old numbers: An old number was already there with probability k/(n-1). It stays in the list if it is not swapped out. Chance of being swapped out = (Chance new number enters) × (Chance this specific spot is picked). Math: (k/n) × (1/k) = 1/n. Chance of surviving = 1 - (1/n) = (n-1)/n. Total Probability: (k/(n-1)) × ((n-1)/n) = k/n. Result: After N items, every item has a k/N probability of being selected. Python Code Here is the implementation using Python's standard library. import random from typing import Iterator, List, Optional def reservoir_sampling(stream: Iterator[int], k: int) -> List[int]: """ Randomly samples k elements from an infinite stream. """ reservoir = [] # Go through the stream one by one for i, item in enumerate(stream): # i starts at 0, so the count is i + 1 current_count = i + 1 if len(reservoir) < k: # If we haven't filled the reservoir yet, add the item reservoir.append(item) else: # We are past the first k elements. # Replace an item in the reservoir with probability k / current_count. # random.random() gives a number between 0.0 and 1.0 if random.random() < k / current_count: # Pick a random spot in the reservoir to replace replace_idx = random.randint(0, k - 1) reservoir[replace_idx] = item return reservoir Part 2: Testing the Solution The Follow-Up Question The interviewer will likely ask: "How do you prove your code is fair? How do I know every number really has an equal chance?" Note: You usually do not need to write a math proof on the whiteboard. Instead, you need to write a simulation script. You run the code many times and check the statistics. Testing Strategy Simulate: Run the reservoir_sampling function many times (e.g., 100,000 times). Count: Track how many times each number gets picked. Compare: Expected: The math says each number should appear Total Trials * (k / N) times. Observed: The actual number of times it appeared. Check: Simple Check: Are the observed numbers close to the expected numbers? Advanced Check: Use a Chi-Square Test. This is a statistical test that tells you if the difference between your results and the expected results is significant or just luck. Python Code for Verification You can use collections.Counter to count and scipy.stats (if allowed) for the math test. from collections import Counter from scipy import stats # Bonus: Shows you know statistical tools def verify_sampling(n: int, k: int, num_trials: int = 100000): """ Checks if Reservoir Sampling is fair. n: Total numbers in the stream (0 to n-1) k: How many to pick num_trials: How many times to run the test """ counts = Counter() for _ in range(num_trials): # Create a fresh stream (0 to n-1) every time stream = iter(range(n)) sample = reservoir_sampling(stream, k) for num in sample: counts[num] += 1 # --- Analyze Results --- print(f"Total elements: {n}, Sample size: {k}, Trials: {num_trials}") expected_count = num_trials * (k / n) print(f"Expected count per element: {expected_count}") # Prepare data for Chi-Square Test observed_frequencies = [] expected_frequencies = [] max_error = 0 print("\n--- Sample Frequencies ---") # Check the counts for each number for i in range(n): obs = counts[i] observed_frequencies.append(obs) expected_frequencies.append(expected_count) # Calculate how far off we are (percentage error) error_pct = abs(obs - expected_count) / expected_count * 100 max_error = max(max_error, error_pct) # Only print a few examples to keep output clean if i < 5 or i > n - 6: print(f"Element {i}: {obs} (Error: {error_pct:.2f}%)") print(f"\nMax Error Percentage: {max_error:.2f}%") # --- Chi-Square Test --- # This tests if the difference is due to randomness or a bad algorithm. # A p-value > 0.05 means the result is likely fair (Uniform Distribution). chi2_stat, p_val = stats.chisquare(f_obs=observed_frequencies, f_exp=expected_frequencies) print("\n--- Chi-Square Test Results ---") print(f"Chi-square statistic: {chi2_stat:.4f}") print(f"P-value: {p_val:.4f}") if p_val > 0.05: print("PASS: P-value > 0.05. The results look fair.") else: print("FAIL: P-value <= 0.05. The results might not be fair.") # Run verification # verify_sampling(n=100, k=5, num_trials=50000) Important Discussion Topics Why use Chi-Square? This test compares what you "observed" (actual counts) against what you "expected" (theoretical counts). If you don't want to use complex statistics, just calculating the "Max Relative Error" (checking if error is under 1-2%) is usually enough for an interview. Memory Usage: The algorithm only saves k items. This means Space Complexity is O(k). However, the verification code saves counts for all N items (O(N)). If N is huge, you might need to test only a small range of numbers to save memory. Randomness Quality: Python's random library is good for simulations but not for security (passwords, cryptography). If you need high security, mention the secrets module. For this interview problem, standard random is fine. Similar Interview Questions Shuffle an Array: (Fisher-Yates shuffle) - Reordering a list randomly. Random Pick Index: (LeetCode 398). Linked List Random Node: (LeetCode 382).