Find Median from Data Stream 解答

Question

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

Solution 1 -- Heap

我们可以依据这个思路来想这道题：

median -> 中位数 -> 中位数两边的subarrays的大小差不超过1 -> 什么数据结构可以满足实时地将输入数组排序，并且保持一半一半的大小？

第一选择是Heap

我们维护一个MaxHeap用来存前半段的小的数据，维护一个MinHeap用来存后半段的大的数据

每次加入一个新数字，我们进行如下顺序判断：

1. MaxHeap是否是空，如果是，加入MaxHeap

2. num是否小于MaxHeap的最大值，如果是，加入MaxHeap

3. MinHeap是否是空，如果是，加入MinHeap

4. MaxHeap和MinHeap都不空

num 和 MaxHeap的最大值 和 MinHeap的最小值 比较

class MedianFinder {
    private PriorityQueue<Integer> minHeap;
    private PriorityQueue<Integer> maxHeap;
    private Double median;
    
    public MedianFinder() {
        minHeap = new PriorityQueue<Integer>(11);
        maxHeap = new PriorityQueue<Integer>(11, Collections.reverseOrder());
        median = null;
    }
    
    // Adds a number into the data structure.
    public void addNum(int num) {
        int size1 = maxHeap.size();
        int size2 = minHeap.size();
        if (size1 == 0) {
            maxHeap.add(num);
        } else if (num <= maxHeap.peek()) {
            maxHeap.add(num);
        } else if (size2 == 0) {
            minHeap.add(num);
        }else {
            int firstMax = maxHeap.peek();
            int secondMin = minHeap.peek();
            if (num <= firstMax) {
                maxHeap.add(num);
            } else if (num >= secondMin) {
                minHeap.add(num);
            } else{
                maxHeap.add(num);
            }
        }
        // Check balance
        size1 = maxHeap.size();
        size2 = minHeap.size();
        if (size1 > size2 + 1) {
            while (size1 > size2 + 1) {
                int top = maxHeap.poll();
                minHeap.offer(top);
                size1--;
                size2++;
            }
        } else if (size2 > size1 + 1) {
            while (size2 > size1 + 1) {
                int top = minHeap.poll();
                maxHeap.offer(top);
                size2--;
                size1++;
            }
        }
        if (size2 == size1 + 1)
            median = (double) minHeap.peek();
        if (size1 == size2 + 1)
            median = (double) maxHeap.peek();
        if (size1 == size2 && size1 > 0)
            median = ((double) maxHeap.peek() + (double) minHeap.peek()) / 2;
    }

    // Returns the median of current data stream
    public double findMedian() {
        return median.doubleValue();
    }
};

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

Solution 2 -- Balanced BST

平衡二叉树也可以满足条件，但是实际代码太多。所以可以在面试时优先用heap写出代码，然后follow-up的时候提出平衡二叉树的思想。