Leetcode 大小堆二分BFPRT二叉排序树AVL数据流的中位数（295）

Posted 2021-03-20 Timeashore

题目

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

[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，Time: O(N), Space: log(N)，居然特么超时了，，，
# class MedianFinder:
#
#     def __init__(self):
#         self.l = []
#
#     def addNum(self, num):
#         self.l.append(num)
#
#     def findMedian(self):
#         length = len(self.l)
#         if length % 2 == 1:
#             return self.findKthLargest(self.l, length // 2 + 1)
#         else:
#             a = self.findKthLargest(self.l, length // 2)
#             b = self.findKthLargest(self.l, length // 2 + 1)
#             return (a + b) / 2
#
#     def findKthLargest(self, nums, k):
#         def getmedian(lis):
#             """返回序列lis中位数，在BFPRT中就是求每5个数小组的中位数"""
#             begin = 0
#             end = len(lis) - 1
#             sum = begin + end
#             mid = sum // 2 + sum % 2  # 这个地方加上sum%2是为了确保偶数个数时我们求的是中间两个数的后一个
#             return sorted(lis)[mid]
#
#         def BFPRT(nums, left, right):
#             """分成每5个数一个小组，并求出每个小组内的中位数"""
#             num = right - left + 1
#             offset = 0 if num % 5 == 0 else 1  # 最后如果剩余的数不足5个，我们也将其分成一个小组，和前面同等对待
#             groups = num // 5 + offset
#             median = []  # 中位数数组
#             for i in range(groups):
#                 begin = left + i * 5
#                 end = begin + 4
#                 Median = getmedian(nums[begin:min(end, right) + 1])
#                 median.append(Median)
#             return getmedian(median)
#
#         def partition(nums, left, right, base):
#             """在 nums[left, right] 将基准base归位"""
#             temp = nums[base]
#             nums[base], nums[right] = nums[right], nums[base]  # 基准和末尾元素互换
#
#             max_index = left
#             for i in range(left, right):  # 把所有小于基准的移到左边
#                 if nums[i] <= temp:  # 要等于啊！这里好坑的说.. 否则通不过[3, 3, 3, 3, 4, 3, 3, 3, 3]  k = 1
#                     nums[max_index], nums[i] = nums[i], nums[max_index]
#                     max_index += 1
#             nums[right], nums[max_index] = nums[max_index], nums[right]  # 基准归位
#             return max_index
#
#         def select(nums, left, right, k_smallest):
#             """在 nums[left, right] 找第k小的元素"""
#             if left == right:  # 递归终止条件
#                 return nums[left]
#             base = BFPRT(nums, left, right)
#             base_index = partition(nums, left, right, nums.index(base))  # 选base为基准，并归位。
#             if base_index == k_smallest:  # 判断目前已归位的基准，是不是第k_smallest位
#                 return nums[k_smallest]
#             elif k_smallest < base_index:  # 递归左半部分
#                 return select(nums, left, base_index - 1, k_smallest)
#             else:  # 递归右半部分
#                 return select(nums, base_index + 1, right, k_smallest)
#
#         return select(nums, 0, len(nums) - 1, len(nums) - k)
#

obj = MedianFinder()
obj.addNum(1)
obj.addNum(2)
print(obj.findMedian())
obj.addNum(3)
print(obj.findMedian())