题目描述
如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。我们使用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/SysuJayceclass 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()