Problem Statement
Design a class to calculate the median of a stream of numbers. The class should have the following two methods:
insertNum(int num): stores the number in the classfindMedian(): returns the median of all numbers inserted in the class
If the count of numbers inserted in the class is even, the median will be the average of the two middle numbers.
Example 1:
1. insertNum(3)
2. insertNum(1)
3. findMedian() -> output: 2.0
4. insertNum(5)
5. findMedian() -> output: 3.0
6. insertNum(4)
7. findMedian() -> output: 3.5
Constraints:
- -10<sup>5</sup> <= num <= 10<sup>5</sup>
- There will be at least one element in the data structure before calling
findMedian. - At most 5 * 10<sup>4</sup> calls will be made to
insertNumandfindMedian.
Why this is a Two Heaps problem
| What the question says | The signal it matches |
|---|---|
| "calculate the median of a stream of numbers" | you need the median of a set |
| "a stream of numbers" | numbers arrive over time |
| "findMedian(): returns the median of all numbers inserted" | the answer is needed after each one |
This is the running median variant: one heap holds the smaller half, the other holds the larger half.
The closest alternative. Keep the numbers in a sorted list. Each insert finds its position by binary search, which is fast, and then shifts everything after it, which is not. The median is then an index.
Count the work before choosing. The constraints allow up to 50,000 calls, so the list can hold 50,000 numbers. Each insert then shifts up to 50,000 of them. That is over a billion moves in the worst case. Two heaps insert in logarithmic time, about 16 steps at this size, and the median is the top of each heap. Sorting the whole list on every call is the same idea done worse. The introduction names it: sorting works, but it is redone after every change.
Solution
As we know, the median is the middle value in an ordered integer list. So a brute force solution could be to maintain a sorted list of all numbers inserted in the class so that we can efficiently return the median whenever required. Inserting a number in a sorted list will take O(N) time if there are N numbers in the list. This insertion will be similar to the Insertion sort. Can we do better than this? Can we utilize the fact that we don't need the fully sorted list - we are only interested in finding the middle element?
Assume x is the median of a list. This means that half of the numbers in the list will be smaller than (or equal to) x and half will be greater than (or equal to) x. This leads us to an approach where we can divide the list into two halves: one half to store all the smaller numbers (let's call it smallNumList) and one half to store the larger numbers (let's call it largeNumList). The median of all the numbers will either be the largest number in the smallNumList or the smallest number in the largeNumList. If the total number of elements is even, the median will be the average of these two numbers.
The best data structure that comes to mind to find the smallest or largest number among a list of numbers is a Heap. Let's see how we can use a heap to find a better algorithm.
- We can store the first half of numbers (i.e.,
smallNumList) in a Max Heap. We should use a Max Heap as we are interested in knowing the largest number in the first half. - We can store the second half of numbers (i.e.,
largeNumList) in a Min Heap, as we are interested in knowing the smallest number in the second half. Inserting a number in a heap will take O(logN), which is better than the brute force approach. At any time, the median of the current list of numbers can be calculated from the top element of the two heaps.
Algorithm Walkthrough
Two heaps hold the numbers, split down the middle. The left heap holds the smaller half with its BIGGEST value on top. The right heap holds the larger half with its SMALLEST value on top. Those two tops are the two numbers either side of the middle, so the median is always right there.
1 of 11
Code
Here is what our algorithm will look like:
Time Complexity
The time complexity of the insertNum() will be O(logN) due to the insertion in the heap. The time complexity of the findMedian() will be O(1) as we can find the median from the top elements of the heaps.
Space Complexity
The space complexity will be O(N) because, as at any time, we will be storing all the numbers.