随时找到数据流中的中位数——Python实现方法

题目：有一个源源不断地吐出整数的数据流，假设你有足够的空间来保存吐出的数。请设计一个名叫MedianHolder的结构，MedianHolder可以随时取得之前吐出所有数的中位数。

解析：关于此问题的主要解题思路为建立大根堆和小根堆，大根堆用来存储较小的数，小根堆用来存储较大的数，在读入数据的过程中要进行大根堆和小根堆的调整，使两者所保存的数据量的差值不大于2，主要的步骤如下：

1. 建立大根堆和小根堆；
2. 读入第一个数据，将其放入大根堆中并建堆；
3. 读入数据，比较该数据与大根堆堆顶元素的大小，如果该值小于大根堆的堆顶，将该值加入大根堆中并进行hepify操作，否则进行下一步操作；
4. 如果小根堆为空，将该值放入小根堆中并建堆，并返回，否则比较该数据与小根堆堆顶元素的大小，如果该值小于小根堆的堆顶元素，则将该值加入到大根堆中，否则将其加入小根堆中，随后对该数据加入的堆进行heapify操作；
5. 加入数据后对，大根堆和小根堆进行相应的调整： 如果大根堆的size比小根堆的size大2，那么从大根堆里将堆顶弹出，并放入小根堆中；如果小根堆的size比大根堆的size大2，那么从小根堆里将堆顶弹出，并放入大根堆中。

建堆完成后，如果大根堆和小根堆的size相等，那么取两个堆得堆顶元素的平均数即可，如果不相等，那么取size较大的堆的堆顶元素即可。

首先，我们利用Python自带的heapq模块分别建立大根堆和小根堆。

import heapq

class BigHeap():
    def __init__(self):
        self.arr = list()
    def heap_insert(self, val):
        heapq.heappush(self.arr, -val)
    def heapify(self):
        heapq.heapify(self.arr)
    def heap_pop(self):
        return -heapq.heappop(self.arr)
    def get_top(self):
        if not self.arr:
            return
        return -self.arr[0]

import heapq

class SmallHeap():
    def __init__(self):
        self.arr = list()
    def heap_insert(self, val):
        heapq.heappush(self.arr, val)
    def heapify(self):
        heapq.heapify(self.arr)
    def heap_pop(self):
        return heapq.heappop(self.arr)
    def get_top(self):
        if not self.arr:
            return
        return self.arr[0]

接下来我们设计MedianHolder结构来获取数据流中的中位数。

class MedianHolder():

    def __init__(self):
        self.bigHeap = BigHeap()
        self.smallHeap = SmallHeap()

    def addNum(self, num):
        if len(self.bigHeap.arr) == 0:
            self.bigHeap.heap_insert(num)
            return
        if self.bigHeap.get_top() >= num:
            self.bigHeap.heap_insert(num)
        else:
            if len(self.smallHeap.arr) == 0:
                self.smallHeap.heap_insert(num)
                return
            if self.smallHeap.get_top() > num:
                self.bigHeap.heap_insert(num)
            else:
                self.smallHeap.heap_insert(num)
        self.modifyTwoHeapSize()

    def getMedian(self):
        smallHeapSize = len(self.smallHeap.arr)
        bigHeapSize = len(self.bigHeap.arr)
        if smallHeapSize + bigHeapSize == 0:
            return None
        smallHeapHead = self.smallHeap.get_top()
        bigHeapHead = self.bigHeap.get_top()
        if (smallHeapSize + bigHeapSize) %2 == 0:
            return (smallHeapHead+bigHeapHead)/2
        else:
            return smallHeapHead if smallHeapSize > bigHeapSize else bigHeapHead

    def modifyTwoHeapSize(self):
        smallHeapSize = len(self.smallHeap.arr)
        bigHeapSize = len(self.bigHeap.arr)
        if smallHeapSize == bigHeapSize + 2:
            self.bigHeap.heap_insert(self.smallHeap.heap_pop())
        if bigHeapSize == smallHeapSize + 2:
            self.smallHeap.heap_insert(self.bigHeap.heap_pop())

测试用例如下：

if __name__ == '__main__':
    arr = [68,51,42,92,13,46,24,58,62,72,32]
    medianHold = MedianHolder()
    for i in range(len(arr)):
        medianHold.addNum(arr[i])
        print(medianHold.getMedian())