对顶堆的应用

题目描述

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

例如，

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

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

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

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

示例 1：

1
2
3
4
输入：
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出：[null,null,null,1.50000,null,2.00000]

示例 2：

1
2
3
4
输入：
["MedianFinder","addNum","findMedian","addNum","findMedian"]
[[],[2],[],[3],[]]
输出：[null,null,2.00000,null,2.50000]

限制：

  • 最多会对 addNum、findMedian 进行 50000 次调用。

分析与求解

是对顶堆的应用，可以用来求动态中位数

  • 维护两个堆，自底向上数值不断变大，下面数值大优先级高，上面数值小优先级高
  • 下面的堆的数量，至多比上面的堆多一

这样就会出现如下的情况：

  • 如果当前是奇数个数，那么肯定是下面的比下面的多
  • 如果是偶数，两个堆的堆顶就是最中间的两个数
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
class MedianFinder {
    PriorityQueue<Integer> down, up;
    int cnt;

    /** initialize your data structure here. */
    public MedianFinder() {
        down = new PriorityQueue<>((a, b) -> b - a);
        up = new PriorityQueue<>((a, b) -> a - b);
        cnt = 0;
    }
    
    public void addNum(int x) {
        if (down.isEmpty() || x <= down.peek()) down.add(x);
        else up.add(x);
        
        if (down.size() > up.size() + 1) up.add(down.poll());
        if (up.size() > down.size()) down.add(up.poll());
        cnt++;
    }
    
    public double findMedian() {
        if ((cnt&1) == 1) return down.peek();
        return (down.peek() + up.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();
 */