【Leetcode 大小堆、二分、BFPRT、二叉排序树、AVL】数据流的中位数（295）
题目
中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。
例如，
[2,3,4] 的中位数是 3
[2,3] 的中位数是 (2 + 3) / 2 = 2.5
设计一个支持以下两种操作的数据结构：
void addNum(int num) - 从数据流中添加一个整数到数据结构中。
double findMedian() - 返回目前所有元素的中位数。
示例：
addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3)
findMedian() -> 2
进阶:
如果数据流中所有整数都在 0 到 100 范围内，你将如何优化你的算法？
如果数据流中 99% 的整数都在 0 到 100 范围内，你将如何优化你的算法？
解答
1，排序——>找中位数，排序用快排，Time: n·log(n)，找中位数复杂度O(1)
2，BFPRT——>没添加一个数，计算一次BFPRT，居然超时了，，，TIme: O(n)
3，二分插入——>找中位数，二分插入平均log(n)，最坏O(n)，找中位数O(1)
4，二叉排序树——>找中位数，构建二叉排序树平均log(n)，最差O(n)，找中位数O(n)
5，AVL平衡二叉树——>优化了二叉排序树最差的情况，让树平衡，Time: log(n)，Space: O(1)，不太好实现
6，大小堆，一个大堆一个小堆，元素先入大堆，再把大堆最大元素给最小堆，如果此时len(小堆)>len(大堆)，则将小堆最小的元素加入大堆，Time: log(n)，Space: O(1)
最优解法：大小堆
通过代码如下：
# # 二分
# from bisect import insort
#
#
# class MedianFinder:
#
# def __init__(self):
# self.l = []
#
# def addNum(self, num):
# insort(self.l, num) # 找到位置log(N)，插入最差时O(N)
#
# def findMedian(self):
# if len(self.l) % 2 != 0:
# return self.l[len(self.l) // 2]
# else:
# return (self.l[len(self.l) // 2] + self.l[len(self.l) // 2 - 1]) / 2
# # 二叉排序树做法，插入操作：Time: 平均log(N)， 最差O(N)，找中位数：O(N)，超时了，，，代码应该是没有问题的，
# # 可以用平衡二叉树(AVL)优化，插入操作：Time: log(N)，中位数在根节点，所以找中位数的复杂度是O(1)
# class BSTstruct:
# # BST的节点结构
#
# def __init__(self, key, left, right):
# self.key = key
# self.left = left
# self.right = right
#
#
# class BST:
# # 二叉排序树类
# def __init__(self):
# self.root = None
# self.start = 0
#
# def insertBST(self, cur_node, key):
# '''
# 构建二叉排序树
# Time: 平均log(N)， 最差O(N)
# '''
# if cur_node == None: # 空树
# self.root = BSTstruct(key, None, None)
#
# elif key <= cur_node.key: # 这里把相同的元素放在了cur_node的左边
# if cur_node.left == None:
# cur_node.left = BSTstruct(key, None, None)
# return
# self.insertBST(cur_node.left, key)
#
# elif key > cur_node.key:
# if cur_node.right == None:
# cur_node.right = BSTstruct(key, None, None)
# return
# self.insertBST(cur_node.right, key)
#
# def findMedian(self, root, cnt):
# """
# 循环着找中位数
# """
# stack = []
# node = root
# cur = 0
# while stack or node:
# while node:
# stack.append(node)
# node = node.left
# node = stack.pop()
# cur += 1
# if cur == cnt:
# return node.key
# node = node.right
#
#
# class MedianFinder:
# def __init__(self):
# self.obj = BST() # init
# self.cnt = 0 # 记录已经插入的个数
#
# def addNum(self, num):
# self.obj.insertBST(self.obj.root, num)
# self.cnt += 1
#
# def findMedian(self):
# '''
# 遍历到第cnt//2个数就是中位数，偶数的话再往下遍历一个
# Time: O(N/2) = O(N)
# '''
# if self.cnt % 2 == 1:
# return self.obj.findMedian(self.obj.root, self.cnt//2+1)
# else:
# a = self.obj.findMedian(self.obj.root, self.cnt//2)
# b = self.obj.findMedian(self.obj.root, self.cnt//2+1)
# return (a+b)/2
# 最优解法：大小堆, Time: log(n), Space: O(1)，优于用二分法维护一个有序列表
# 两个堆，一个大堆存一半小元素，一个小堆存一半大元素
# 遍历，先入大堆，把大堆最大的给小堆，如果小堆个数大于大堆，就再把小堆最小给大堆
# 最终，奇数时，大堆比小堆多一个元素，堆顶就是中位数；偶数时，取两堆顶计算即可
from heapq import heappop, heappush
class MedianFinder:
def __init__(self):
self.small = []
self.large = []
def addNum(self, num):
heappush(self.large, -num)
t = -heappop(self.large)
heappush(self.small, t)
if len(self.small) > len(self.large):
t = heappop(self.small)
heappush(self.large, -t)
def findMedian(self):
if len(self.small) < len(self.large):
return -self.large[0]
else:
return (self.small[0]-self.large[0])/2
# 这题有多种解法
# BFPRT解法 (Timeout...)
# Time: O(n), Space: O(n)
class MedianFinder:
def __init__(self):
self.nums = []
self.len = 0
def addNum(self, num: int) -> None:
self.nums.append(num)
self.len += 1
def findMedian(self) -> float:
if self.len % 2 == 0:
return (self.BFPRT(self.nums, 0, self.len-1, self.len//2) + self.BFPRT(self.nums, 0, self.len-1, self.len//2-1)) / 2
return self.BFPRT(self.nums, 0, self.len-1, self.len//2)
def BFPRT(self, nums, left, right, K):
"""BFPRT算法"""
if left == right:
return nums[left] # 不是排序，这要返回值
base = self.medianOfMedians(nums, left, right) # 找基准
base_index = self.partition(nums, left, right, nums.index(base)) # 基准归位
if base_index == K:
return nums[base_index]
elif base_index > K:
return self.BFPRT(nums, left, base_index - 1, K) # 递归左边。不是快排，这要返回值
else:
return self.BFPRT(nums, base_index + 1, right, K) # 递归右边
def medianOfMedians(self, nums, left, right):
"""找中位数基准"""
length = right - left + 1
offset = 0 if length % 5 == 0 else 1 # 最后不够5个的算一组
groups = length // 5 + offset
medians = []
for i in range(groups):
start = left + i * 5
end = start + 4
medians.append(self.get_median(nums[start: min(end, right) + 1]))
return self.BFPRT(medians, 0, groups - 1,
groups // 2) # 这里递归BFPRT，保证得到的是“准确中位数”（30% ~ 70%）; 而非“近似中位数” 防止时间复杂度退化
def get_median(self, nums):
"""找5个数的中位数"""
return sorted(nums)[len(nums) // 2] # 常数级别
def partition(self, nums, left, right, base):
"""常规partition"""
if left >= right:
return
temp = nums[base]
nums[base], nums[right] = nums[right], nums[base]
max_index = left
for i in range(left, right):
if nums[i] <= temp:
nums[i], nums[max_index] = nums[max_index], nums[i]
max_index += 1
nums[max_index], nums[right] = nums[right], nums[max_index]
return max_index
# Your MedianFinder object will be instantiated and called as such:
obj = MedianFinder()
obj.addNum(1)
obj.addNum(2)
print(obj.findMedian()) # -> 1.5
obj.addNum(3)
print(obj.findMedian()) # -> 2
