Find Median from Data Stream

hardFREEHeap / Priority QueueBlind 75

Design a data structure that supports adding integers and finding the median efficiently.

Implement a class MedianFinder with:

add_num(num): add the integer num to the stream.
find_median(): return the median of all numbers added so far as a number (not a string). With an even count, the median is the mean of the two middle values.

Input: operations as add,num or find_median separated by semicolons. The test harness parses them, calls your methods, and prints each median result with 1 decimal place.

Example

Input

add,1;add,2;find_median;add,3;find_median

Output

1.5
2.0

Reasoning

The input is a sequence of operations read left to right: add,1 and add,2 put 1 and 2 into the stream, find_median asks for the median of everything added so far, then add,3 grows the stream before the second query.
The median is the middle value once the numbers are in sorted order. At the first query the stream holds $\{1, 2\}$ , an even count, so there are two middle values and the median is their mean: $\frac{1+2}{2} = 1.5$ .
At the second query the stream holds $\{1, 2, 3\}$ , an odd count, so the single middle value (the 2nd of 3) is the median: $2$ , shown with one decimal place as 2.0.
The output has one line per find_median query, in order: 1.5 then 2.0. An earlier answer is never revised when later numbers arrive.
Notice that insertion order is irrelevant: adding 3, 1, 2 gives the same sorted stream $1, 2, 3$ and the same median, $2.0$ .

Constraints

-10^5 <= num <= 10^5
At least one add before any median call

Learn this in

Heaps & Priority QueuesDSA Study Plan

More from Blind 75

Alien Dictionaryhard Clone Graphmedium Combination Summedium Word Search IIhard Merge K Sorted Listshard

Concept primer

Background Knowledge

The problem of finding the median from a data stream is a classic example of an online algorithm, where we need to process the data in a sequential manner and make decisions based on the current state of the system. To solve this problem efficiently, we need to understand the concept of heaps and priority queues. A heap is a specialized tree-based data structure that satisfies the heap property: the parent node is either greater than (or less than) its child nodes. This property allows us to efficiently extract the minimum or maximum element from the heap.

The median of a dataset is the middle value when the data is sorted in ascending order. If the dataset has an even number of elements, the median is the average of the two middle values. To find the median efficiently, we can use two heaps: a max heap to store the smaller half of the numbers and a min heap to store the larger half. By maintaining the balance between the two heaps, we can ensure that the median is always accessible in constant time.

The key concept here is to understand how to maintain the balance between the two heaps. We need to ensure that the size of the max heap is always greater than or equal to the size of the min heap, and the maximum element in the max heap is less than or equal to the minimum element in the min heap. This balance is crucial in maintaining the correctness of the median calculation.

Algorithm/Approach

The general approach to solve this problem is to use a combination of heaps and priority queues to maintain the balance between the smaller and larger halves of the numbers. We can use the following algorithm pattern:

Use two heaps: a max heap and a min heap
When adding a new number, determine which heap to add it to based on the current state of the heaps
After adding the number, rebalance the heaps to maintain the correct balance
When finding the median, use the top elements of the two heaps to calculate the median

Step-by-Step Strategy

To implement the solution, follow these steps:

Initialize two empty heaps: a max heap and a min heap
When adding a new number, compare it with the top element of the max heap (if it exists)
If the new number is less than or equal to the top element of the max heap, add it to the max heap
Otherwise, add it to the min heap
After adding the number, check if the max heap has more than one more element than the min heap
If so, remove the top element from the max heap and add it to the min heap
When finding the median, check if the total number of elements is odd or even
If it's odd, the median is the top element of the max heap
If it's even, the median is the average of the top elements of the max heap and the min heap

Common Pitfalls

When implementing the solution, watch out for the following common pitfalls:

Not maintaining the correct balance between the two heaps
Not handling the case where the total number of elements is odd or even correctly
Not using the correct data structure (e.g., using a single heap instead of two)

Time & Space Complexity

The expected time complexity for adding a number is $O(\log n)$ , where $n$ is the total number of elements. The expected time complexity for finding the median is $O(1)$ . The space complexity is $O(n)$ , where $n$ is the total number of elements.

solution.py

Python 3.13

Test results