劍指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()