Data Structures & Algorithms in Java – Heaps – Find Median from Data Stream

Problem:

Given a input stream of integers find the median of them.

For example:

The input has three methods , you need to implement all of them:

MedianFinder() – this is the constructor , called once in the beginning

addNum() – this method keeps adding a single number at a time to the stream

findMedian() – this method finds the median of the given numbers.

Median is the middlemost number in a sorted array if the number of numbers is odd.

If the number of numbers is even , then median is the average of the middlemost two numbers after the numbers are sorted.

Here is a sample input:

["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]

And the output:

[null, null, null, 1.5, null, 2.0]

Try out the solution here:

https://leetcode.com/problems/find-median-from-data-stream/

Solution:

Hint:

The brute force solution to this problem is to sort the input numbers in a list and then find the median

Sorting takes O(nlogn) time.

A better solution exists using heaps.

Explanation:

To find the median you can do the following:

Heap would be a great data structure to implement this for the following reasons:

All you have to do is keep the heap size same as you keep adding each element.

Also you need to push all the bigger elements to the min heap and the smaller elements to the max heap.

You can do this balancing by following the below algorithm:

ALGORITHM:

This way we can make sure the first heap (max heap) contains the sorted first half of the input

And the second heap(min heap) contains the sorted second half of the input.

In a real world use case ,

Imagine you have two bags.

And you need to put stones of different sizes into these bags.

One bag is for smaller ones.

Other is for bigger ones .

You want equal number of small and big stones .

One exception is if the total number of stones is odd you can have one more in the first bag of smaller stones.

You take the first stone first and put it in the small bag

Since there is only one stone there is no other stone to compare with, so let it be in the small bag

Now pick the second stone and put it in the small bag again.

Of the two stones pick the largest and put it in the big bag

Now the small bag has one small stone and the big bag has one big stone , exactly what we want

Now pick the third stone and put it in the small bag

Take the biggest of the two stones and put it in the big bag

Now the big bag size is greater

So pick the smallest stone from the big bag and put it in the small bag

Now the smallest stones are in the small bag

The biggest stones are in the big bag

Continue this process for all the stones.

You end up with smaller stones on the small bag and bigger ones in the big bag

Here is the code:

class MedianFinder {
    
    Queue<Integer> maxHeap;
    
    Queue<Integer> minHeap;
    
    public MedianFinder() {
    
        maxHeap = new PriorityQueue<Integer>((a,b)->{
            
            
            return Integer.compare(b,a);
        });
        
        
        minHeap = new PriorityQueue<Integer>();
        
     
    }
    
    public void addNum(int num) {
        
        
      
        maxHeap.add(num);
        
        
        minHeap.add(maxHeap.poll());
        
        
        if(maxHeap.size() < minHeap.size()){
            
            
            maxHeap.add(minHeap.poll());
        }
        
    }
    
    public double findMedian() {
      
      return maxHeap.size() > minHeap.size()? maxHeap.peek(): (double)(maxHeap.peek() + minHeap.peek()) /2; 
        
    }
}

Time complexity:

There are five heap operations here :

Each heap operation takes O(logn) time.

So time complexity above is O(5logn).

Also to find the median takes O(1) time (constant irrespective of the input size)

So total time complexity is approximately O(logn)

Space complexity:

We use two heaps each of them holding half of the input.

So space complexity is O(n)

That’s it!