双堆法求数据流的中位数

在这里插入图片描述

代码逻辑

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

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

1.添加元素

添加新元素时，如果大顶堆为空,或新元素 ≤ 大顶堆堆顶 → 放入大顶堆，否则 → 放入小顶堆。如果两个堆的大小差为 2,就从元素多的堆中取出堆顶,放入另一个堆。

2.查找中位数

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

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

代码实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42


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());
            }
        }
    }
}