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