双堆法求数据流的中位数

代码逻辑

这道题的关键在于如何在动态添加元素的过程中,快速找到中位数。

我们可以用一个大顶堆 (maxHeap)存储较小的一半元素，堆顶是这部分的最大值 和一个小顶堆 (minHeap)存储较大的一半元素,堆顶是这部分的最小值，这样的话，中位数就在两个堆的堆顶附近!

1.添加元素

添加新元素时，如果大顶堆为空,或新元素 ≤ 大顶堆堆顶 → 放入大顶堆，否则 → 放入小顶堆。

如果两个堆的大小差为 2,就从元素多的堆中取出堆顶,放入另一个堆。

2.查找中位数

如果两个堆大小相等 → 中位数 = (大顶堆堆顶 + 小顶堆堆顶) / 2;

如果大小不等 → 中位数 = 元素多的那个堆的堆顶。

代码实现

java 复制代码

class MedianFinder {
    private PriorityQueue<Integer> maxHeap;  // 大顶堆,存储较小的一半
    private PriorityQueue<Integer> minHeap;  // 小顶堆,存储较大的一半
    
    public MedianFinder() {
        // Java的PriorityQueue默认是小顶堆,需要自定义比较器实现大顶堆
        maxHeap = new PriorityQueue<>((a, b) -> b - a);
        minHeap = new PriorityQueue<>((a, b) -> a - b);
    }
    
    public void addNum(int num) {
        // 决定新元素放入哪个堆
        if (maxHeap.isEmpty() || maxHeap.peek() >= num) {
            maxHeap.add(num);
        } else {
            minHeap.add(num);
        }
        // 调整两个堆的平衡
        balance();
    }
    
    public double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            // 偶数个元素,返回两个堆顶的平均值
            return (double) (maxHeap.peek() + minHeap.peek()) / 2;
        } else {
            // 奇数个元素,返回元素多的那个堆的堆顶
            return maxHeap.size() > minHeap.size() ? maxHeap.peek() : minHeap.peek();
        }
    }
    
    private void balance() {
        // 如果两个堆的大小差为2,需要调整
        if (Math.abs(maxHeap.size() - minHeap.size()) == 2) {
            if (maxHeap.size() > minHeap.size()) {
                minHeap.add(maxHeap.poll());
            } else {
                maxHeap.add(minHeap.poll());
            }
        }
    }
}