Sunday, January 3, 2016

[LeetCode] Find Median from Data Stream

思路是使用一个minHeap和maxHeap。不知道有没有更好的方法。
class MedianFinder { 

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

    public MedianFinder() {
        minHeap = new PriorityQueue<Integer>();
        maxHeap = new PriorityQueue<Integer>(500, Collections.reverseOrder());
    }


    // Adds a number into the data structure.
    public void addNum(int num) {
        if (maxHeap.size() == 0 || num <= maxHeap.peek()) {
            maxHeap.add(num);
        }
        else {
            minHeap.add(num);
        }
        if (maxHeap.size() - minHeap.size() >= 2) {
            minHeap.add(maxHeap.poll());
        }
        else if (minHeap.size() > maxHeap.size()) {
            maxHeap.add(minHeap.poll());
        }
    }


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

No comments:

Post a Comment