Python 中的大顶堆和小顶堆 (heapq库)

## 定义
大顶堆：一种经过排序的完全二叉树，其中任意非终端节点数值*均不小于其左子节点和右子节点的值*。如果一颗二叉树满足最小堆的要求，那么，堆顶（根节点）也就是整个序列的最大元素。
* **arr[i] >= arr[2i+1] && arr[i] >= arr[2i+2]**
小顶堆：也是一种经过排序的完全二叉树，其中任意非终端节点数值*均不大于其左子节点和右子节点的值*。如果一棵二叉树满足最小堆的要求，那么，堆顶（根节点）也就是整个序列的最小元素。
* **arr[i] <= arr[2i+1] && arr[i] <= arr[2i+2]**
![最大堆/最小堆示例](https://zoe.red/usr/uploads/2024/02/3042802495.png)
python中(小顶堆)堆是通过*数组*来实现的，其中的元素从 0 开始计数，对于所有的 *k* 都有 `a[k]<=a[2*k+1]` 且 `a[k]<=a[2*k+2]`。 为了便于比较，不存在的元素被视为无穷大。 堆最有趣的特性在于 `a[0]` 总是其中最小的元素。
* <span style='color:red'>任意升序排列的数组都可以看做是小顶堆/小根堆</span>
上面的特殊*不变量*是用来作为一场锦标赛的高效内存表示。 下面的数字是索引 *k* 而不是 `a[k]`:
```bash
              1                                 2
      3               4                5               6
  7       8       9       10      11      12      13      14
15 16   17 18   19 20   21 22   23 24   25 26   27 28   29 30
## heapq库常用方法
**注意：** python中的heapq库只有小顶堆，没有大顶堆，当使用大顶堆时，可以在插入元素时将元素取反，弹出是也取反，一些常用操作如下：
```python
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.查询最大/最小的n个前n个元素
# heapq.nsmallest() 查询堆中的最小n个元素
# heapq.nlargest() 查询堆中的最大n个元素
# 注意: 输入为随机数组即可, 无需有序, 操作前后内容不变
n = 3
heap = [1,3,5,7,2,4,6,8]
print(heapq.nsmallest(n,heap)) # [1,2,3]
print(heapq.nlargest(n,heap)) # [8,7,6]
*补充1: 堆排序的自定义比较函数*
heapq内置的方法并不支持自定义比较方法, 但是可以灵活变通, 示例如下
* F1: 使用Python 内置模块 functools 的高阶函数`cmp_to_key`, 通过`cmp_to_key(obj_cmp_function)(obj)`将原生不支持比较的对象转换为可以比较的，其中`obj_cmp_function`自定义比较函数输入参数为2个, 返回值类型为<span style='color:red'>小于则返回一个负数，相等则返回零，大于则返回一个正数</span>
  * 另外, 在一些接受key的函数中（例如sorted，min，max，heapq.nlargest，itertools.groupby），key仅仅支持一个参数，就无法实现两个及以上参数之间的对比，采用cmp\_to\_key 函数，可以接受两个参数，将两个参数做处理，比如做和做差
* F2: 对类对象进行扩展包装, 扩展对象内实现 `__lt__`方法
<div class="tab-container post_tab box-shadow-wrap-lg">
<ul class="nav no-padder b-b scroll-hide" role="tablist">
<li class='nav-item active' role="presentation"><a class='nav-link active' style="" data-toggle="tab" 
aria-controls='tabs-f0c02a598240c86d075bbf804f451cbf650' role="tab" data-target='#tabs-f0c02a598240c86d075bbf804f451cbf650'>F2</a></li><li class='nav-item ' role="presentation"><a class='nav-link ' style="" data-toggle="tab" 
aria-controls='tabs-1160fcbd7c51d0a8086c17a483f3844d661' role="tab" data-target='#tabs-1160fcbd7c51d0a8086c17a483f3844d661'>F1</a></li>
<div class="tab-content no-border">
<div role="tabpanel" id='tabs-f0c02a598240c86d075bbf804f451cbf650' class="tab-pane fade active in">
```python
# 类对象扩展包装, 扩展对象原生支持__lt__方法
# 合并k个升序链表代码
# leetcode链接: https://leetcode.cn/problems/merge-k-sorted-lists/
# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
# 为ListNode定制的包含__lt__方法的扩展类
class ListNodeExt:
    def __init__(self, node):
        self.node = node 
    def __lt__(self, other):
        return self.node.val < other.node.val
class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        import heapq
        # 排序函数
        # heapq不直接支持比较函数, 弃用此方案
        # cmp_node = lambda x: x.val
        node_list_heap = []
        for node in lists:
            if node is not None:
                node_list_heap += [ListNodeExt(node)] 
        new_node_list = []
        # 构建堆
        heapq.heapify(node_list_heap)
        while node_list_heap:
            # 将最小的/根 node 弹出, 记录并更新索引
            small_node_ext = heapq.heappop(node_list_heap)
            node = small_node_ext.node
            new_node_list += [node]
            # 弹出节点后, 不为空的话重新压入下一个指向节点到堆中
            if node.next is not None:
                heapq.heappush(node_list_heap, ListNodeExt(node.next))
        for index in range(len(new_node_list)-1):
            new_node_list[index].next = new_node_list[index+1]
        if not new_node_list:
            return None 
        else:
            return new_node_list[0]
</div><div role="tabpanel" id='tabs-1160fcbd7c51d0a8086c17a483f3844d661' class="tab-pane fade  ">
```python
# 合并k个升序链表代码
# leetcode链接: https://leetcode.cn/problems/merge-k-sorted-lists/
# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        import heapq
        from functools import cmp_to_key
        # 排序函数
        # heapq不直接支持比较函数
        # 修改, 使用系统提供的比较封装库函数, 使得对象可以比较
        # 注意: 当前的示例返回是相减的结果,
        cmp_node = cmp_to_key(lambda x, y: x.val - y.val)
        node_list_heap = []
        for node in lists:
            if node is not None:
                # node_list_heap += [ListNodeExt(node)] 
                node_list_heap += [cmp_node(node)] 
        new_node_list = []
        # 构建堆
        heapq.heapify(node_list_heap)
        while node_list_heap:
            # 将最小的/根 node 弹出, 记录并更新索引
            small_node_ext = heapq.heappop(node_list_heap)
            node = small_node_ext.obj
            new_node_list += [node]
            # 弹出节点后, 不为空的话重新压入下一个指向节点到堆中
            if node.next is not None:
                heapq.heappush(node_list_heap, cmp_node(node.next))
        for index in range(len(new_node_list)-1):
            new_node_list[index].next = new_node_list[index+1]
        if not new_node_list:
            return None 
        else:
            return new_node_list[0]
# 扩展示例
# 1. python原生, 简单排序
lst = [(9, 4), (2, 10), (4, 3), (3, 6)]
lst.sort(key=lambda item: item[0])
print(lst)
# 2. 复杂排序:
# 修改规则:
# * 如果元组里第一个元素是奇数，就用元组里第一个元素进行排序
# * 如果元组里第一个元素是偶数，则用这个元组里的第二个元组进行大小比较，
# 面对这样的需求，列表的sort方法无法满足, 此时使用cmp_to_key
from functools import cmp_to_key
lst = [(9, 4), (2, 10), (4, 3), (3, 6)]
def cmp(x, y):
    a = x[0] if x[0] %2 == 1 else x[1]
    b = y[0] if y[0] %2 == 1 else y[1]
    return 1 if a > b else -1 if a < b else 0
lst.sort(key=cmp_to_key(cmp))
print(lst)
# 3. cmp_to_key实现实例
# 在内部定义了一个类K， 并使用传入的cmp函数完成了比较关系运算符的重载，函数返回的是一个类
# 因为在python中，类和函数都是callable的，这里把类当成了函数来用完全可行
def cmp_to_key(mycmp):
    """Convert a cmp= function into a key= function"""
    class K(object):
        __slots__ = ['obj']
        def __init__(self, obj):
            self.obj = obj
        def __lt__(self, other):
            return mycmp(self.obj, other.obj) < 0
        def __gt__(self, other):
            return mycmp(self.obj, other.obj) > 0
        def __eq__(self, other):
            return mycmp(self.obj, other.obj) == 0
        def __le__(self, other):
            return mycmp(self.obj, other.obj) <= 0
        def __ge__(self, other):
            return mycmp(self.obj, other.obj) >= 0
        __hash__ = None
    return K
*补充:2 heapq提供的大顶堆API*
| 小根堆/最小堆         | 小根堆/最小堆                |
| --------------------- | ---------------------------- |
| **heapq.heappush**    | 无                           |
| **heapq.heappop**     | **heapq.\_heappop\_max**     |
| **heapq.heapify**     | **heapq.\_heapify\_max**     |
| **heapq.heapreplace** | **heapq.\_heapreplace\_max** |
源码中没有 `heappush` 对应的 `_heappush_max` 方法，具体 大根堆/最大堆 怎么用，大家还是自己摸索吧~ 😂
上面 大根堆/最大堆 的 API 是用 Python 实现的，也可以添加如下 import 语句，使用 C 实现的 API ~
```python
from _heapq import _heapify_max as heapify
from _heapq import _heappop_max as heappop
from _heapq import _heapreplace_max as heapreplace
## 小试牛刀
题目1：[剑指 Offer 40. 最小的k个数](https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/ "剑指 Offer 40. 最小的k个数")
```python
# 最大堆解法，参考题解：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://leetcode-cn.com/problems/shu-ju-liu-zhong-de-zhong-wei-shu-lcof/ "剑指 Offer 41. 数据流中的中位数")
```python]
# 最大堆最小堆解法，参考题解: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]
题目3: [Leetcode: 数组中的第K个最大元素](https://leetcode.cn/problems/kth-largest-element-in-an-array/)
* 建堆之后, 弹出元素k的结果就是
* 时间复杂度: 建堆$O(n)$, 弹出k次调整复杂度 $O(k*log(n))$
* 建堆的过程, 需要特别留意要从第一个非叶子节点开始往前遍历: (max_len -1 )//2
* 调整的过程, 需要从当前节点开始, 与其它子节点比较, 若无需交换, 即可终止; 如需要交换, 发生交换的子节点需要重复以上过程
* 要能自己手写完成建堆与调整函数, 特别是其中的maxHeapify函数
```python
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # 自己实现大根堆的建堆与弹出调整
        # from heapq import _heapify_max as heapify
	# 堆调整, 建立在i节点以下部分已经是堆的前提下
        def maxHeapify(heap, i):
            # 将索引位置i处节点下部分完全大根堆
            # 要求 i 节点以下部分已经完全符合最大堆定义
            max_len = len(heap)
            while i < max_len:
                l = 2*i+1
                r = 2*i+2
                if l >= max_len: # 无子节点, 更新完成
                    break
                largest = l 
                if r < max_len and heap[l] < heap[r]:
                    largest = r 
                if heap[largest] > heap[i]:
                    heap[i], heap[largest] = heap[largest], heap[i]
                    i = largest 
                else: # 无需交换即退出/终止!!!
                    break
        def buildMaxHeap(heap):
            # 从第一个非叶子节点开始
            max_len = len(heap)
            for i in range((max_len-1)//2, -1, -1):
                maxHeapify(heap, i)
        heapify = buildMaxHeap
        def pop(heap):
            data = heap[0]
            heap[0] = heap[-1]
            heap.pop() # 末尾移除
            maxHeapify(heap, 0)
            return data
        heapify(nums)
        k_largest = None
        for i in range(k):
            k_largest = pop(nums)
        return k_largest
## 参考:
* https://docs.python.org/zh-cn/3/library/heapq.html
* https://leetcode.cn/circle/discuss/zM9yCy/
* https://www.gairuo.com/p/python-cmp-to-key
* http://www.coolpython.net/python_senior/standard_module/functools_cmp_to_key.html
* https://stackoverflow.com/questions/7803121/in-python-heapq-heapify-doesnt-take-cmp-or-key-functions-as-arguments-like-sor
* https://blog.csdn.net/wangws_sb/article/details/124108070