對應 LeetCode 295 資料流中的中位數

問題描述

設計一個支援如下兩種操作的資料結構：

• void addNum(int)：從資料流中獲取一個元素，新增到當前的資料結構中

• double findMedian()：返回當前資料結構中儲存的資料的中位數

解決思路

由於這裡無法確切知道資料元素的規模，因此一般通過列表的方式儲存元素再求取中位數的方式不是特別可靠。

考慮使用 “堆” 資料結構來完成這個功能，維護兩個堆：最大堆和最小堆，最大堆用於儲存中位數元素以下的所有元素，最小堆儲存大於等於中位數的所有元素。

為了方便，將最小堆中的堆頂元素視為奇數大小的資料流的中位數，在新增時注意小心地調整這兩個堆中的元素以使得滿足上面的條件

實現

class MedianFinder {
    final PriorityQueue<Integer> minPq; // 維護最小堆，儲存大於等於中位數的元素
    final PriorityQueue<Integer> maxPq; // 維護最大堆，儲存小於中位數的元素

    public MedianFinder() {
        minPq = new PriorityQueue<>((x, y) -> x - y);
        maxPq = new PriorityQueue<>((x, y) -> y - x);
    }
    
    public void addNum(int num) {
        /*
            注意這裡堆中元素的調整
        */
        if (minPq.size() != maxPq.size()) {
            minPq.offer(num);
            maxPq.offer(minPq.poll());
        } else {
            maxPq.offer(num);
            minPq.offer(maxPq.poll());
        }
    }
    
    public double findMedian() {
        /*
            由於這兩個堆滿足了我們給定的條件，因此中位數的計算就變得簡單
        */
        if (minPq.size() != maxPq.size()) 
            return minPq.peek();

        return (minPq.peek() + maxPq.peek()) * 1.0 / 2;
    }
}
 

複雜度分析：

• 時間複雜度：由於每次呼叫 addNum(int) 方法時都會觸發堆的平衡操作，因此時間複雜度為 \(O(log_2n)\)，對於中位數的計算，由於只是獲取了兩個堆的堆頂元素，因此時間複雜度為 \(O(1)\)

• 空間複雜度：需要額外的空間來儲存輸入的資料，因此空間複雜度為 \(O(n)\)

參考：

^[1] https://leetcode-cn.com/problems/find-median-from-data-stream/

^[2] https://leetcode-cn.com/problems/shu-ju-liu-zhong-de-zhong-wei-shu-lcof/solution/mian-shi-ti-41-shu-ju-liu-zhong-de-zhong-wei-shu-y/