数据流中的中位数 Find Median from Data Stream

2019-04-17 16:34:50

问题描述：

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构：

• void addNum(int num) - 从数据流中添加一个整数到数据结构中。
• double findMedian() - 返回目前所有元素的中位数。

示例：

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

进阶:

1. 如果数据流中所有整数都在 0 到 100 范围内，你将如何优化你的算法？
2. 如果数据流中 99% 的整数都在 0 到 100 范围内，你将如何优化你的算法？

问题求解：

class MedianFinder {
    PriorityQueue<Integer> minheap;
    PriorityQueue<Integer> maxheap;

    /** initialize your data structure here. */
    public MedianFinder() {
        minheap = new PriorityQueue<Integer>((Integer o1, Integer o2) -> o1.compareTo(o2));
        maxheap = new PriorityQueue<Integer>((Integer o1, Integer o2) -> o2.compareTo(o1));
    }

    public void addNum(int num) {
        if (maxheap.isEmpty() || maxheap.peek() >= num) maxheap.add(num);
        else minheap.add(num);
        while (minheap.size() > maxheap.size()) maxheap.add(minheap.poll());
        while (maxheap.size() > minheap.size() + 1) minheap.add(maxheap.poll());
    }

    public double findMedian() {
        int size = minheap.size() + maxheap.size();
        if (size % 2 != 0) return 1.0 * maxheap.peek();
        else return 0.5 * minheap.peek() + 0.5 * maxheap.peek();
    }
}