Coding questions · Amazon
Design a streaming median structure
This interview question evaluates algorithm design, data structures, correctness, complexity, edge cases, and implementation details in a realistic interview setting. A strong answer for Design a streaming median structure states assumptions, handles edge cases, explains trade-offs, and shows how to validate the result clearly.
Design a data structure for a number stream that supports: add(num: int) in O(log n) time and median() -> float in O(1) or O(log n) time. It must handle up to 10^6 operations, duplicates, and both even/odd counts. Describe the APIs, internal invariants, and prove the time/space complexity. Implement the core methods.
What a Strong Answer Covers
- A clear algorithm with the right data structures and enough pseudocode or code-level detail to implement it.
- A correctness argument that explains why the algorithm covers all required cases.
- Time and space complexity, plus at least one alternative approach when relevant.
- Focused tests for normal cases, edge cases, and failure modes.
Follow-up Questions
- How would the approach change if the input were streaming or too large for memory?
- What invariants would you assert in production code?
- Which tests would catch off-by-one, duplicate, or tie-breaking bugs?
Streaming Median Structure (Find Median from Data Stream)
Implement a data structure for a number stream that supports two operations and runs each operation efficiently enough to handle up to 10^6 operations with duplicates and both even/odd counts:
- add(num): insert an integer into the stream in O(log n) time.
- median(): return the median of all numbers added so far in O(1) time. If the count is even, the median is the average of the two middle values; if odd, it is the single middle value. The median of an empty stream is undefined (return None / null).
Driver format: you are given a list of operations to replay in order. Each operation is a list:
- ["add", num] -> insert num, contributes None to the output.
- ["median"] -> append the current median (a float, or None if the stream is empty). Return a list with exactly one entry per operation.
Internal invariants (two-heap design): keep a max-heap lo for the smaller half and a min-heap hi for the larger half, maintaining len(lo) == len(hi) or len(lo) == len(hi) + 1. The median is the top of lo when sizes differ, or the average of both tops when sizes are equal. add() pushes onto the correct heap then rebalances by at most one move (O(log n)); median() reads heap tops (O(1)). Space is O(n).
Constraints
- 1 <= number of operations <= 10^6
- -10^9 <= num <= 10^9
- median() may be called on an empty stream (return None / null)
- Duplicates are allowed; both even and odd total counts must be handled
Examples
Example 1
Input: ([['add', 1], ['add', 2], ['median'], ['add', 3], ['median']],)
Output: [None, None, 1.5, None, 2.0]
Explanation: After 1,2 the median is (1+2)/2 = 1.5 (even count). After adding 3, the median of {1,2,3} is 2.0 (odd count).
Example 2
Input: ([['median']],)
Output: [None]
Explanation: Median of an empty stream is undefined, so None is returned.
Hints
- Maintain two heaps: a max-heap for the smaller half and a min-heap for the larger half, keeping their sizes balanced to within one element.
- On add(), push to the correct half based on the current max-heap top, then move one element across if the size invariant breaks — this keeps the operation O(log n).
- On median(), if the heaps are equal size average the two tops; otherwise the larger heap's top is the median. This reads in O(1).
Solution
Solution Alignment
The prompt asks for an implementation-level answer. The safest way to present it is to define the state, maintain clear invariants, then walk through complexity and tests.
Problem Restatement
Design a data structure for a number stream that supports: add(num: int) in O(log n) time and median() -> float in O(1) or O(log n) time. It must handle up to 10^6 operations, duplicates, and both even/odd counts. Describe the APIs, internal invariants, and prove the time/space complexity. Implement the core methods.
Recommended Approach
Start with a brute-force baseline to confirm correctness, then identify the repeated work or ordering property that enables a better data structure such as a hash map, heap, stack, queue, two pointers, prefix sums, BFS/DFS, or dynamic programming. Write the implementation around a small invariant and test that invariant directly.
Correctness
The implementation should maintain an invariant after each loop or operation that directly matches the problem statement. At termination, that invariant implies the returned value has considered every valid candidate exactly once, or has preserved the required data-structure state after every API call.
Complexity
State the baseline complexity and the optimized complexity. For most interview constraints, justify why the optimized approach meets the expected input size.
Edge Cases and Tests
Empty and singleton inputs, duplicates, ties, invalid inputs, boundary values, and tests that exercise the main invariant.