劍指Offer_#41_ 資料流中的中位數

Contents

• • 題目
  • 思路分析
  • 解答

題目

如何得到一個數據流中的中位數？如果從資料流中讀出奇數個數值，那麼中位數就是所有數值排序之後位於中間的數值。如果從資料流中讀出偶數個數值，那麼中位數就是所有數值排序之後中間兩個數的平均值。
例如，
[2,3,4]的中位數是 3
[2,3] 的中位數是 (2 + 3) / 2 = 2.5
設計一個支援以下兩種操作的資料結構：
void addNum(int num) - 從資料流中新增一個整數到資料結構中。
double findMedian() - 返回目前所有元素的中位數。
示例 1：

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

示例 2：

輸入：
["MedianFinder","addNum","findMedian","addNum","findMedian"]
[[],[2],[],[3],[]]
輸出：[null,null,2.00000,null,2.50000]

限制：
最多會對addNum、findMedia進行50000次呼叫。

思路分析

藉助Java中的優先佇列PriorityQueue,分別建立大頂堆和小頂堆，大頂堆儲存資料流中較小的一半數字，小頂堆儲存資料流中較大的一半數字。如下圖所示。

解答

class MedianFinder
 
 {
    Queue<Integer> A,B;
    /** initialize your data structure here. */
    public MedianFinder() {
        A = new PriorityQueue<>();//小頂堆，儲存較大的一半
        B = new PriorityQueue<>((x,y) -> y-x);//大頂堆，儲存較小的一半
    }
    
    public void addNum(int num) {
        if(A.size() == B.size()){
            //先加入A，再加入B,如果長度相同，說明上次加入的是B,這次該加入A
            //先把num加入B,再把B的頂彈出加入到A(即把B中最大的加入A)
            B.add(num);
            A.add(B.poll());
        }else{
            //同理，先把num加入A,再把A的頂彈出加入到B(把A中的最小加入到B)
            A.add(num);
            B.add(A.poll());
        }
    }
    
    public double findMedian() {
        //整個資料流長度為偶數
        if(A.size() == B.size()) return (A.peek() + B.peek()) / 2.0;
        //整個資料流長度為奇數，那麼中位數就是A的頂
        else return A.peek();
    }
}