LeetCode 295. 数据流的中位数

原题链接

题目描述

思路

代码

C++

class MedianFinder {

    priority_queue<int, vector<int>, greater<int>> up;
    priority_queue<int> down;

public:
    /** initialize your data structure here. */
    MedianFinder() {

    }
    
    void addNum(int num) {
        if ( down.empty() || num <= down.top() )
        {
            down.push(num);
            if ( down.size() > up.size() + 1 )
            {
                up.push(down.top());
                down.pop();
            }
        }
        else
        {
            up.push(num);
            if ( up.size() > down.size() )
            {
                down.push(up.top());
                up.pop();
            }
        }
    }
    
    double findMedian() {
        if ( down.size() + up.size() & 1 ) return down.top();
        return (down.top() + up.top()) / 2.0;
    }
};

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

Java

class MedianFinder {

    private PriorityQueue<Integer> up;
    private PriorityQueue<Integer> down;

    /** initialize your data structure here. */
    public MedianFinder() {
        up = new PriorityQueue<>();
        down = new PriorityQueue<>((a, b) -> b - a);
    }
    
    public void addNum(int num) {
        if ( down.isEmpty() || num <= down.peek() )
        {
            down.offer(num);
            if ( down.size() > up.size() + 1 )
                up.offer(down.poll());
        }
        else
        {
            up.offer(num);
            if ( up.size() > down.size() )
                down.offer(up.poll());
        }
    }
    
    public double findMedian() {
        if ( (up.size() + down.size()) % 2 == 1 ) return down.peek();
        return (up.peek() + down.peek()) / 2.0;
    }
}

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