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

題目描述
如何獲得一個數據流中的中位數？若是從數據流中讀出奇數個數值，那麼中位數就是全部數值排序以後位於中間的數值。若是從數據流中讀出偶數個數值，那麼中位數就是全部數值排序以後中間兩個數的平均值。咱們使用Insert()方法讀取數據流，使用GetMedian()方法獲取當前讀取數據的中位數。

# -*- coding: utf-8 -*-
# @Time         : 2019-07-09 10:29
# @Author       : Jayce Wong
# @ProjectName  : job
# @FileName     : getMedianFromStream.py
# @Blog         : https://blog.51cto.com/jayce1111
# @Github       : https://github.com/SysuJayce

class Solution:
    """
    要求一個數據流中的中位數，因爲要求的中位數隨着數據流的改變也會發生變化，所以最樸素的解法就是在
    輸入數據的時候直接用一個數組保存起來，而後在須要返回中位數的時候先對數組進行排序而後計算中位數。

    可是若是要求中位數，咱們其實並不須要獲得一個徹底有序的數組。若是數組的元素個數是奇數，那麼中位
    數就是排序後的中間元素，若是是偶數個，中位數就是排序後中間兩個元素的平均值。基於這個觀察，若是
    咱們可以維護兩個指針p1, p2，分別指向當前數組的左右兩部分，其中p1指向的是左邊部分的最大值，p2
    指向的是右邊部分的最小值。若是當前數組個數是奇數個，那麼p1和p2指向同一個位置。

    要實現上述的兩個指針，咱們能夠利用堆排序中的知識。
    p1至關於指向最大堆的堆頂，p2至關於指向最小堆的堆頂。
    """
    def __init__(self):
        self.min_heap = []
        self.max_heap = []

    def adjustHeap(self, data, idx):
        """
        對於一個堆，當咱們對其進行調整的時候，首先要找到給定節點的子節點，而後判斷子節點是否符合
        最大堆/最小堆的要求，若是不符合那就調整。
        :param data: 待調整的堆
        :param idx: 待調整的節點下標
        """
        length = len(data)  # 當前堆的大小
        temp = data[idx]  # 先記錄待調整節點的值，最後再放到適當的位置中
        k = 2 * idx + 1  # 左子節點的下標
        while k < length:  # 咱們的調整須要在樹內調整，所以須要限定下標
            # 這裏因爲咱們有兩個堆，而最大堆和最小堆的條件不同，因此設置這個分支
            if id(data) == id(self.max_heap):
                # 在最大堆的調整中，咱們只須要找到左右子節點中最大的值而後跟根節點進行調整
                if k + 1 < length and data[k + 1] > data[k]:
                    k += 1
                # 這裏用賦值而非交換，由於待調整節點可能最終並不位於這個位置，而咱們以前已經記錄
                # 好了待調整節點的值，所以能夠放心賦值
                if data[k] > temp:
                    data[idx] = data[k]
                    idx = k  # 賦值完以後須要記錄最新的待調整節點可能位於的位置
                # 一旦出現子樹符合堆的定義就能夠終止調整，由於咱們是從下往上構造堆的，只要當前
                # 子樹符合定義，那麼再往下的子樹也符合定義
                else:
                    break
            elif id(data) == id(self.min_heap):
                if k + 1 < length and data[k + 1] < data[k]:
                    k += 1
                if data[k] < temp:
                    data[idx] = data[k]
                    idx = k
                else:
                    break
            # 調整完當前子樹以後再調整孫子節點
            k = 2 * k + 1
        # 最後要記得將待調整節點的值放到正確的位置
        data[idx] = temp

    def Insert(self, num):
        # 若是已經有偶數個元素，那麼這第奇數個插入到最小堆(右邊)
        if (len(self.min_heap) + len(self.max_heap)) & 1 == 0:
            # 在插入最小堆以前，需確認待插入元素比最大堆的全部元素都大(即比最大堆的堆頂元素大)
            # 不然須要先插入最大堆，而後將調整後的最大堆的堆頂挪到最小堆中
            if self.max_heap and num < self.max_heap[0]:
                self.max_heap.append(num)
                num = self.max_heap.pop(0)
                self.adjustHeap(self.max_heap, 0)

            # 插入到最小堆後要調整獲得新的堆頂。
            # 因爲咱們是從0開始構造堆的，因此只須要對堆頂進行調整便可
            self.min_heap.append(num)
            self.adjustHeap(self.min_heap, 0)
        # 若是已經有奇數個元素，那麼這第偶數個元素插入到最大堆(左邊)
        else:
            if self.min_heap and num > self.min_heap[0]:
                self.min_heap.append(num)
                num = self.min_heap.pop(0)
                self.adjustHeap(self.min_heap, 0)

            self.max_heap.append(num)
            self.adjustHeap(self.max_heap, 0)

    def GetMedian(self, num):
        # 這裏num參數是牛客網的OJ有問題，只有這樣才能成功調用GetMedian()
        if not self.min_heap:
            raise Exception("No numbers are available.")

        if (len(self.min_heap) + len(self.max_heap)) & 1 == 0:
            return (self.min_heap[0] + self.max_heap[0]) / 2.0
        else:
            return self.min_heap[0]

def main():
    solution = Solution()
    nums = [5, 2, 3, 4, 1, 6, 7, 0, 8]
    for num in nums:
        solution.Insert(num)
        print(solution.GetMedian(num))

if __name__ == '__main__':
    main()