[LeetCode]Find Median from Data Stream

580 查看

Find Median from Data Stream

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

Examples:
[2,3,4] , the median is 3

[2,3], the median is (2 + 3) / 2 = 2.5

Design a data structure that supports the following two operations:

  • void addNum(int num) - Add a integer number from the data stream to the data structure.

  • double findMedian() - Return the median of all elements so far.

For example:

add(1)
add(2)
findMedian() -> 1.5
add(3) 
findMedian() -> 2

分析

这道题的经典做法就是维护两个Heap, 一个MaxHeap, 一个MinHeap,维护他们的size,比如保证

MaxHeap.size() - MinHeap.size() >= 1

当然这种题可以有些follow up, 比如:

  • 除了heap做,还有什么其他方法?
    BST可以,比如可以在Node里增加一个size表示整个children的数量,findMedian时间复杂度会增加到log(n);

  • 如果是找比如处在1/10th的元素,怎么找?
    同样可以用两个heap做,维护一个的size是另一个size的1/9即可。

复杂度

time: addNum -> O(log(n)), findMedian -> O(1)
space: O(n)

代码

class MedianFinder {

    // Adds a number into the data structure.
    PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>();
    PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(Collections.reverseOrder());
    
    public void addNum(int num) {
        maxHeap.add(num);
        minHeap.add(maxHeap.remove());
        
        // 维护两个heap的size 
        if (minHeap.size() > maxHeap.size()) 
            maxHeap.add(minHeap.remove());
    }

    // Returns the median of current data stream
    public double findMedian() {
        if (maxHeap.size() == minHeap.size()) 
            return (maxHeap.peek() + minHeap.peek()) / 2.0;
        else 
            return maxHeap.peek();
    }
};

// Your MedianFinder object will be instantiated and called as such:
// MedianFinder mf = new MedianFinder();
// mf.addNum(1);
// mf.findMedian();