[劍指offer] 資料流中的中位數

本文首發于我的個人部落格： 尾尾部落

題目描述

如何得到一個資料流中的中位數？如果從資料流中讀出奇數個數值，那麼中位數就是所有數值排序之後位于中間的數值。如果從資料流中讀出偶數個數值，那麼中位數就是所有數值排序之後中間兩個數的平均值。我們使用Insert()方法讀取資料流，使用GetMedian()方法擷取目前讀取資料的中位數。

解題思路

我們可以将資料排序後分為兩部分，左邊部分的資料總是比右邊的資料小。那麼，我們就可以用最大堆和最小堆來裝載這些資料：

• 最大堆裝左邊的資料，取出堆頂（最大的數）的時間複雜度是O(1)
• 最小堆裝右邊的資料，同樣，取出堆頂（最小的數）的時間複雜度是O(1)

從資料流中拿到一個數後，先按順序插入堆中：如果左邊的最大堆是否為空或者該數小于等于最大堆頂的數，則把它插入最大堆，否則插入最小堆。然後，我們要保證左邊的最大堆的size等于右邊的最小堆的size或者最大堆的size比最小堆的size大1。

要擷取中位數的話，直接判斷最大堆和最小堆的size，如果相等，則分别取出兩個堆的堆頂除以2得到中位數，不然，就是最大堆的size要比最小堆的size大，這時直接取出最大堆的堆頂就是我們要的中位數。

參考代碼

import java.util.PriorityQueue;
import java.util.Comparator;
public class Solution {
    // 最小堆（右）
    private PriorityQueue<Integer> rHeap = new PriorityQueue<>(); 
    // 最大堆（左）
    private PriorityQueue<Integer> lHeap = new PriorityQueue<Integer>(15, new Comparator<Integer>() {
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
    // 保證lHeap.size()>=rHeap.size()
    public void Insert(Integer num) {
        // 先按大小插入，再調整
        if(lHeap.isEmpty() || num <= lHeap.peek())
            lHeap.offer(num);
        else
            rHeap.offer(num);
        
        if(lHeap.size() < rHeap.size()){
            lHeap.offer(rHeap.peek());
            rHeap.poll();
        }else if(lHeap.size() - rHeap.size() == 2){
            rHeap.offer(lHeap.peek());
            lHeap.poll();
        }
    }
    public Double GetMedian() {
        if(lHeap.size() > rHeap.size())
            return new Double(lHeap.peek());
        else
            return new Double(lHeap.peek() + rHeap.peek())/2;
    }
}