【劍指Offer系列41】數據流中的中位數

題目

如何得到一個數據流中的中位數？如果從數據流中讀出奇數個數值，那麼中位數就是所有數值排序之後位於中間的數值。如果從數據流中讀出偶數個數值，那麼中位數就是所有數值排序之後中間兩個數的平均值。

例如，

[2,3,4] 的中位數是 3

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

設計一個支持以下兩種操作的數據結構：

void addNum(int num) - 從數據流中添加一個整數到數據結構中。
double findMedian() - 返回目前所有元素的中位數。

代碼

Python

from heapq import *

class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.small_heap = [] # 最小堆，存儲較大的一半
        self.big_heap = [] # 最大堆，存儲較小的一半

    def addNum(self, num: int) -> None:
        # 最大堆的元素都小於最小堆的元素，但num不一定比最小堆中元素小，故先將num插入最小堆，再將最小堆堆頂元素插入最大堆
        if len(self.small_heap) != len(self.big_heap): 
            heappush(self.small_heap, num)
            heappush(self.big_heap, -heappop(self.small_heap))
        else:
            heappush(self.big_heap, -num)
            heappush(self.small_heap, -heappop(self.big_heap))

    def findMedian(self) -> float:
        return self.small_heap[0] if len(self.small_heap) != len(self.big_heap) else (self.small_heap[0] - self.big_heap[0])/2.0


# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()

C++

class MedianFinder {
    
    priority_queue<int,vector<int>,greater<int>> small_heap;
    priority_queue<int,vector<int>,less<int>> big_heap;
    int count;
    
public:
    /** initialize your data structure here. */
    MedianFinder() {

    }
    
    void addNum(int num) {
        if (count%2 != 0) {
            small_heap.push(num);
            big_heap.push(small_heap.top());
            small_heap.pop();
        }
        else {
            big_heap.push(num);
            small_heap.push(big_heap.top());
            big_heap.pop();
        }
        count++;
    }
    
    double findMedian() {
        if (count%2==0) {
            return (small_heap.top() + big_heap.top())/2.0;
        }
        else {
            return (double)(small_heap.top());
        }
    }
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder* obj = new MedianFinder();
 * obj->addNum(num);
 * double param_2 = obj->findMedian();
 */