劍指offer：數(shù)據(jù)流中的中位數(shù)

題目描述
如何得到一個(gè)數(shù)據(jù)流中的中位數(shù)？如果從數(shù)據(jù)流中讀出奇數(shù)個(gè)數(shù)值，那么中位數(shù)就是所有數(shù)值排序之后位于中間的數(shù)值。如果從數(shù)據(jù)流中讀出偶數(shù)個(gè)數(shù)值，那么中位數(shù)就是所有數(shù)值排序之后中間兩個(gè)數(shù)的平均值。我們使用Insert()方法讀取數(shù)據(jù)流，使用GetMedian()方法獲取當(dāng)前讀取數(shù)據(jù)的中位數(shù)。

# -*- 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:
    """
    要求一個(gè)數(shù)據(jù)流中的中位數(shù)，由于要求的中位數(shù)隨著數(shù)據(jù)流的改變也會(huì)發(fā)生變化，因此最樸素的解法就是在
    輸入數(shù)據(jù)的時(shí)候直接用一個(gè)數(shù)組保存起來，然后在需要返回中位數(shù)的時(shí)候先對(duì)數(shù)組進(jìn)行排序然后計(jì)算中位數(shù)。

    但是如果要求中位數(shù)，我們其實(shí)并不需要得到一個(gè)完全有序的數(shù)組。如果數(shù)組的元素個(gè)數(shù)是奇數(shù)，那么中位
    數(shù)就是排序后的中間元素，如果是偶數(shù)個(gè)，中位數(shù)就是排序后中間兩個(gè)元素的平均值?；谶@個(gè)觀察，如果
    我們能夠維護(hù)兩個(gè)指針p1, p2，分別指向當(dāng)前數(shù)組的左右兩部分，其中p1指向的是左邊部分的最大值，p2
    指向的是右邊部分的最小值。如果當(dāng)前數(shù)組個(gè)數(shù)是奇數(shù)個(gè)，那么p1和p2指向同一個(gè)位置。

    要實(shí)現(xiàn)上述的兩個(gè)指針，我們可以利用堆排序中的知識(shí)。
    p1相當(dāng)于指向最大堆的堆頂，p2相當(dāng)于指向最小堆的堆頂。
    """
    def __init__(self):
        self.min_heap = []
        self.max_heap = []

    def adjustHeap(self, data, idx):
        """
        對(duì)于一個(gè)堆，當(dāng)我們對(duì)其進(jìn)行調(diào)整的時(shí)候，首先要找到給定節(jié)點(diǎn)的子節(jié)點(diǎn)，然后判斷子節(jié)點(diǎn)是否符合
        最大堆/最小堆的要求，如果不符合那就調(diào)整。
        :param data: 待調(diào)整的堆
        :param idx: 待調(diào)整的節(jié)點(diǎn)下標(biāo)
        """
        length = len(data)  # 當(dāng)前堆的大小
        temp = data[idx]  # 先記錄待調(diào)整節(jié)點(diǎn)的值，最后再放到適當(dāng)?shù)奈恢弥?        k = 2 * idx + 1  # 左子節(jié)點(diǎn)的下標(biāo)
        while k < length:  # 我們的調(diào)整需要在樹內(nèi)調(diào)整，因此需要限定下標(biāo)
            # 這里由于我們有兩個(gè)堆，而最大堆和最小堆的條件不一樣，所以設(shè)置這個(gè)分支
            if id(data) == id(self.max_heap):
                # 在最大堆的調(diào)整中，我們只需要找到左右子節(jié)點(diǎn)中最大的值然后跟根節(jié)點(diǎn)進(jìn)行調(diào)整
                if k + 1 < length and data[k + 1] > data[k]:
                    k += 1
                # 這里用賦值而非交換，因?yàn)榇{(diào)整節(jié)點(diǎn)可能最終并不位于這個(gè)位置，而我們之前已經(jīng)記錄
                # 好了待調(diào)整節(jié)點(diǎn)的值，因此可以放心賦值
                if data[k] > temp:
                    data[idx] = data[k]
                    idx = k  # 賦值完之后需要記錄最新的待調(diào)整節(jié)點(diǎn)可能位于的位置
                # 一旦出現(xiàn)子樹符合堆的定義就可以終止調(diào)整，因?yàn)槲覀兪菑南峦蠘?gòu)造堆的，只要當(dāng)前
                # 子樹符合定義，那么再往下的子樹也符合定義
                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
            # 調(diào)整完當(dāng)前子樹之后再調(diào)整孫子節(jié)點(diǎn)
            k = 2 * k + 1
        # 最后要記得將待調(diào)整節(jié)點(diǎn)的值放到正確的位置
        data[idx] = temp

    def Insert(self, num):
        # 如果已經(jīng)有偶數(shù)個(gè)元素，那么這第奇數(shù)個(gè)插入到最小堆(右邊)
        if (len(self.min_heap) + len(self.max_heap)) & 1 == 0:
            # 在插入最小堆之前，需確認(rèn)待插入元素比最大堆的所有元素都大(即比最大堆的堆頂元素大)
            # 否則需要先插入最大堆，然后將調(diào)整后的最大堆的堆頂挪到最小堆中
            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)

            # 插入到最小堆后要調(diào)整得到新的堆頂。
            # 由于我們是從0開始構(gòu)造堆的，所以只需要對(duì)堆頂進(jìn)行調(diào)整即可
            self.min_heap.append(num)
            self.adjustHeap(self.min_heap, 0)
        # 如果已經(jīng)有奇數(shù)個(gè)元素，那么這第偶數(shù)個(gè)元素插入到最大堆(左邊)
        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參數(shù)是?？途W(wǎng)的OJ有問題，只有這樣才能成功調(diào)用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()