LeetCode Hard: Find Median from Data Stream

LeetCode Problem #295: Find Median from Data Stream
Link to Problem Description: https://leetcode.com/problems/find-median-from-data-stream/
Category: Applications of Binary Heaps (or Priority Queues)
Solution Approach: Divide elements into 2 mutually-balanced Priority Queues - one having elements smaller than the middle element(s), one having larger.
Time Complexity of Approach:

• O(lg n) for insertion where n = number of elements seen so far.
• O(n) for finding the median of elements seen so far. Space Complexity: O(n) 

 

This LeetCode Hard Problem is a pretty neat application of Binary Heaps, which are readily usable as PriorityQueues in Java.

I have discussed the possible approaches to dolve the problem in the adjoining youtube video link and in this blog post I present 2 solutions for the Twin-heap based solution:

 

Approach #1: Using PriorityQueue Collection Class of Java Collections Framework

 

 

class MedianFinder {

    PriorityQueue lesser;
    PriorityQueue greater;
    
    public MedianFinder() {
        lesser = new PriorityQueue<>(10, (o1, o2)->o2-o1);
        greater = new PriorityQueue<>();
    }
    
    public void addNum(int num) { // O(lg n)
        if (lesser.size() == 0 || num < lesser.peek()) {
            lesser.add(num);
        } else {
            greater.add(num);
        }
        
        //Balance the heaps
        if (Math.abs(lesser.size() - greater.size()) <= 1)
            return;
        
        if (lesser.size() < greater.size())
            lesser.add(greater.poll());
        else
            greater.add(lesser.poll());
    }
    
    public double findMedian() { // O(1)
        if (lesser.size() == greater.size())
            return ((double)lesser.peek() + greater.peek())/2;
        
        if (lesser.size() > greater.size())
            return lesser.peek();
        return greater.peek();
    }
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */

 

 

Approach 2: Without using pre-build classes of Java - To test if you can code a Binary Heap data structure from scratch!

 

 

class MyPriorityQueue> {
    private ArrayList q;
    private Comparator comp;
    
    public MyPriorityQueue(int initialCapacity, Comparator cmp) {
        q = new ArrayList(initialCapacity);
        comp = cmp;
    }
    public MyPriorityQueue() {
        q = new ArrayList(10);
        comp = (o1, o2)->o1.compareTo(o2);
    }
    
    private int parent(int idx) {
        int pos = idx+1;
        return pos/2 - 1;
    }
    private int left(int idx) {
        int pos = idx+1;
        return (pos*2) -1;
    }
    private int right(int idx) {
        int pos = idx+1;
        return (pos*2+1) -1;
    }
    
    public void add (E e) {
        q.add(e);
        
        //bubble-up if necessary
        int currentIdx = q.size()-1;
        int parentIdx = parent(currentIdx);
        while (parentIdx >= 0 && comp.compare(e, q.get(parentIdx)) < 0) {
            q.set(currentIdx, q.get(parentIdx));
            q.set(parentIdx, e);
            currentIdx = parentIdx;
            parentIdx = parent(currentIdx);
        }
    }
    
    public int size() {
        return q.size();
    }
    
    public E peek() {
        if (q.isEmpty())
            return null;
        return q.get(0);
    }
    
    public void remove() {
        if (q.isEmpty())
            return;
        q.set(0, q.get(q.size()-1));
        q.remove(q.size()-1);
        minHeapify(0);
    }
    
    private void minHeapify(int i) {
        if (i >= q.size())
            return;
        
        int leastValIdx = i;
        if (left(i) < q.size() && comp.compare(q.get(leastValIdx), q.get(left(i))) > 0)
            leastValIdx = left(i);
        if (right(i) < q.size() && comp.compare(q.get(leastValIdx), q.get(right(i))) > 0)
            leastValIdx = right(i);
        
        if (leastValIdx != i) {
            E tmp = q.get(leastValIdx);
            q.set(leastValIdx, q.get(i));
            q.set(i, tmp);
            
            minHeapify(leastValIdx);
        }
    }
}

class MedianFinder {

    /** initialize your data structure here. */
    MyPriorityQueue lesserElementsMaxHeap;
    MyPriorityQueue greaterElementsMinHeap;
    
    public MedianFinder() {
        lesserElementsMaxHeap = new MyPriorityQueue<>(10, (o1, o2)->o2-o1);
        greaterElementsMinHeap = new MyPriorityQueue<>(); //minHeap by default
    }
    
    public void addNum(int num) {
        Integer lesserMax = lesserElementsMaxHeap.peek();
        if (lesserMax == null || lesserMax>num)
            lesserElementsMaxHeap.add(num);
        else
            greaterElementsMinHeap.add(num);
        
        
        //ensure balanced heaps
        if (Math.abs(lesserElementsMaxHeap.size() - greaterElementsMinHeap.size()) <= 1)
            return;
        
        if (lesserElementsMaxHeap.size() > greaterElementsMinHeap.size()) {
            lesserMax = lesserElementsMaxHeap.peek();
            greaterElementsMinHeap.add(lesserMax);
            lesserElementsMaxHeap.remove();
        } else {
            Integer greaterMin = greaterElementsMinHeap.peek();
            lesserElementsMaxHeap.add(greaterMin);
            greaterElementsMinHeap.remove();
        }
            
    }
    
    public double findMedian() {
        if (lesserElementsMaxHeap.size() == greaterElementsMinHeap.size())
            return ((double)lesserElementsMaxHeap.peek() + greaterElementsMinHeap.peek())/2;
        
        if (lesserElementsMaxHeap.size() > greaterElementsMinHeap.size())
            return lesserElementsMaxHeap.peek();
        else
            return greaterElementsMinHeap.peek();
    }
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */

 

 

 Video explanation for the approach:

 

 

Share your thoughts in comments!

 

<3

Become a Patron!