October 11, 2019

295. Find Median from Data Stream

295. Find Median from Data Stream

Refer to this post

左半边用maxHeap,右半边用minHeap。
maxHeap holds the smaller elements of the stream with the ability to provide the largest element in it in O(1)
minHeap holds the larger elements of the stream with the ability to provide the least element in it in O(1)

Time complexity: O(logn) for insert into heap, O (1) for findMedian. Space complexity: O(n)

class MedianFinder {

    PriorityQueue<Integer> maxHeap;
    PriorityQueue<Integer> minHeap;

    /** initialize your data structure here. */
    public MedianFinder() {
        maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
        minHeap = new PriorityQueue<>();
    }

    public void addNum(int num) {
        // 1. Insert value
        // Left size >= right size
        if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
            maxHeap.offer(num);
        } else {
            minHeap.offer(num);
        }

        // 2. Balance left/right
        if (maxHeap.size() < minHeap.size()) {
            maxHeap.offer(minHeap.poll());
        } else if (maxHeap.size() - minHeap.size() == 2) {
            minHeap.offer(maxHeap.poll());
        }
    }

    public double findMedian() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.peek();
        } else {
            return (minHeap.peek() + maxHeap.peek()) / 2.0;
        }
    }

}
comments powered by Disqus