劍指offer | 數據流中的中位數

數據流中的中位數

描述

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

代碼 (Java)

import java.util.Comparator;
import java.util.PriorityQueue;
public class Solution {
    // 最小堆 （保存較大的一半數字）
    public PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>();

    // 最大堆 （保存較小的一半數字
    public PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(15, new Comparator<Integer>() {
        @Override
        public int compare(Integer o1, Integer o2) {
            return o2 - o1;
        }
    });

    int count = 0;      // 記錄總數目（奇偶性）

    public void Insert(Integer num) {
        // 奇數時，插入最小堆
        if ((count & 1) == 0) {
            // 如果新的數據比最大堆的一些數小，則要把最大堆的堆頂元素插入最小堆
            if (maxHeap.size() > 0 && num < maxHeap.peek()) {
                maxHeap.offer(num);
                num = maxHeap.poll();   // 改爲最大堆的最大元素
            }
            minHeap.offer(num);
        } else {    // 偶數時，插入最大堆
            // 如果新的數據比最小堆的一些數大，則要把最小堆的堆頂元素插入最大堆
            if (minHeap.size() > 0 && num > minHeap.peek()) {
                minHeap.offer(num);
                num = minHeap.poll();   // 改爲最小堆的最小元素
            }
            maxHeap.offer(num);
        }
        count++;
    }

    public Double GetMedian() {
        if ((count & 1) == 0)
            return new Double((minHeap.peek() + maxHeap.peek())) / 2;
        else 
            return new Double(minHeap.peek());
    }
}