python 中的最大堆和最小堆(heapq库)

目录

首先来看一下什么是最大堆和最小堆？

python heapq库中的一些常用方法

小试牛刀

---

首先来看一下什么是最大堆和最小堆？

最大堆：一种经过排序的完全二叉树，其中任意非终端节点数值均不小于其左子节点和右子节点的值。如果一颗二叉树满足最小堆的要求，那么，堆顶（根节点）也就是整个序列的最大元素。

最小堆：也是一种经过排序的完全二叉树，其中任意非终端节点数值均不大于其左子节点和右子节点的值。如果一棵二叉树满足最小堆的要求，那么，堆顶（根节点）也就是整个序列的最小元素。

python heapq库中的一些常用方法

注意：python中的heapq库只有最小堆，没有最大堆，当使用最大堆时，可以在插入元素时将元素取反，弹出是也取反，一些常用操作如下：

import heapq
# [2,0,4,1]

# 1.创建堆
# 方法一：定义一个空列表，然后使用heapq.heqppush(item)函数把元素加入到堆中
item = 2
heap = []
heapq.heappush(heap,item)
# 方法二：使用heapq.heapify(list)将列表转换为堆结构
heap = [2,0,4,1]
heapq.heapify(heap)

# 2.heapq.heappush() 添加新元素 num
num = 3
heapq.heappush(heap,num)

# 3.heapq.heappop() 删除并返回堆顶元素
heapq.heappop(heap)

# 4.heapq.heappushpop() 比较添加元素num与堆顶元素的大小:如果num>堆顶元素，删除并返回堆顶元素，然后添加新元素num;如果num<堆顶元素，返回num，原堆不变
# 其实也就等价于 添加新元素num，然后删除并返回堆顶元素
num = 0
heapq.heappushpop(heap,num)

# 5.heapq.heapreplace() 删除并返回堆顶元素，然后添加新元素num
num = 5
heapq.heapreplace(heap,num)

# 6. heapq.merge() 合并多个排序后的序列成一个排序后的序列， 返回排序后的值的迭代器。
heap1 = [1,3,5,7]
heap2 = [2,4,6,8]
heap = heapq.merge(heap1,heap2)
print(list(heap))

# 7.heapq.nsmallest() 查询堆中的最小n个元素
n = 3
heap = [1,3,5,7,2,4,6,8]
print(heapq.nsmallest(n,heap)) # [1,2,3]

# 8.heapq.nlargest() 查询堆中的最大n个元素
n = 3
heap = [1,3,5,7,2,4,6,8]
print(heapq.nlargest(n,heap)) # [8,7,6]

小试牛刀

题目1：剑指 Offer 40. 最小的k个数

# 最大堆解法，参考题解：https://www.bilibili.com/video/BV1To4y1d7cW/
import heapq
class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return []
        # python中只有最小堆，没有最大堆
        # 将所有元素取反，弹出的时候也取反
        heap = [-x for x in arr[0:k]]
        heapq.heapify(heap)
        for x in arr[k:]:
            if -x > heap[0]:
                heapq.heapreplace(heap,-x)
        return [-x for x in heap]

题目2：剑指 Offer 41. 数据流中的中位数

# 最大堆最小堆解法，参考题解:https://www.bilibili.com/video/BV1J5411J7yj/
import heapq
class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.maxHeap = [] # 最大堆保存的是最小的n/2个数
        self.minHeap = [] # 最小堆保存的是最大的n/2个数

    def addNum(self, num: int) -> None:
        # 若两堆的数目相等，就让minHeap的元素个数+1
        # 具体做法分为两步 1.将新元素加入maxHeap 2.将maxHeap的堆顶元素加入minHeap
        if len(self.maxHeap)==len(self.minHeap):
            heapq.heappush(self.minHeap,-heapq.heappushpop(self.maxHeap,-num))
        else:
            heapq.heappush(self.maxHeap,-heapq.heappushpop(self.minHeap,num))
            
    def findMedian(self) -> float:
        if len(self.maxHeap)==len(self.minHeap):
            return (self.minHeap[0]-self.maxHeap[0])/2.0
        else:
            return self.minHeap[0]