Java: Find Median from Data Stream

Thought Process

To efficiently find the median of a stream of numbers, we need a dynamic data structure approach. We use two heaps: a max-heap for the lower half and a min-heap for the upper half. The key is to maintain balanced heaps where their sizes differ by at most 1. When adding a number, we first push it to one heap and then rebalance by moving the top element to the other heap. The median is either the top of the larger heap or the average of the top elements when sizes are equal. 

class MedianFinder {

    PriorityQueue<Integer> firstHalf;
    PriorityQueue<Integer> secondHalf;

    public MedianFinder() {
        firstHalf = new PriorityQueue<>((a,b)->(b-a)); // Max Heap
        secondHalf = new PriorityQueue<>((a,b)->(a-b)); // Min Heap
    }
    
    public void addNum(int num) {
        
        int diffInSize = Math.abs(firstHalf.size()-secondHalf.size());
        // Both half has equal elements
        if(diffInSize==0){
            // Atlast, we have to put smallest element from the secondHalf into firstHalf
            secondHalf.offer(num);
            firstHalf.offer(secondHalf.poll());
        }
        else if(diffInSize==1){
            // Atlast, we have to put largest element from the firstHalf into secondHalf
            firstHalf.offer(num);
            secondHalf.offer(firstHalf.poll());
        }
    }
    
    public double findMedian() {

        int diffInSize = Math.abs(firstHalf.size()-secondHalf.size());
        if(diffInSize==0)
            return (firstHalf.peek() + secondHalf.peek())/2.0;
        else
            return firstHalf.peek();

    }
}

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

Code Complexity

Time Complexity: O(log n) per insertion

Each insertion involves heap operations which take logarithmic time relative to the number of elements.

Space Complexity: O(n)

We store all elements across both heaps, requiring linear space proportional to the input size.

Code copied!