【题目描述】
如何得到一个数据流中的中位数?如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。我们使用Insert()方法读取数据流,使用GetMedian()方法获取当前读取数据的中位数。
【解题思路】
举一个例子:比如1,2,3,4,5,6,7,8的中位数是4和5的平均值,4是前一半的最大值,5是后一半的最小值。这给我们的启示是,可以用优先队列来解这道题,维护一个大根堆和一个小根堆,使用大根堆存储比较小的那一半数据,使用小根堆存储比较大的那一半数据。只要这两个堆是平衡的(两个堆的元素数量相等或相差1),则每次求中位数就只需要访问大根堆堆顶的元素和小根堆堆顶的元素——如果是奇数个数,大根堆堆顶元素就是中位数;如果是偶数个数,大根堆堆顶和小根堆堆顶元素的平均值就是中位数。
具体实现过程为:
- Insert():每次读入一个数的时候,先将其插入到大根堆中,然后再将大根堆堆顶的元素(较小的一半数据中最大的元素)插入到小根堆中;如果此时有奇数个数,将小根堆堆顶的元素(较大的一半数据中最小的元素)插入到大根堆中(这一步是为了实现大根堆和小根堆的平衡)。
- GetMedian():如果是奇数个数,大根堆堆顶元素就是中位数;如果是偶数个数,大根堆堆顶和小根堆堆顶元素的平均值就是中位数。
有一点需要注意的是优先队列的主要操作是插入和删除最小元素(元素的值本身为优先级键值,小元素享有高优先级)。所以大根堆的操作就要对元素相应地取反。
用Python实现的代码如下:
import heapq
class Solution:
def __init__(self):
self.count = 0
self.max_heap = []
self.min_heap = []
def Insert(self, num):
# write code here
self.count += 1
heapq.heappush(self.max_heap, -num)
max_heap_top = heapq.heappop(self.max_heap)
heapq.heappush(self.min_heap, -max_heap_top)
if self.count & 1:
min_heap_top = heapq.heappop(self.min_heap)
heapq.heappush(self.max_heap, -min_heap_top)
def GetMedian(self,s):
# write code here
if self.count & 1:
return -self.max_heap[0]
else:
return (self.min_heap[0] - self.max_heap[0]) / 2.0