堆 heap

堆是一种特殊的树形数据结构，它通常是一个完全二叉树。堆分为两种类型：最大堆和最小堆。（heapq库中的堆默认是最小堆）

  • 最大堆：每个节点的值都大于或等于其子节点的值。
  • 最小堆：每个节点的值都小于或等于其子节点的值。

heapq库中常用函数

  1. heapq.heapify()

    heapify 函数用于将一个列表转换为堆结构。在 Python 中，默认使用的是最小堆。该函数会对列表进行原地操作，即直接修改原列表，而不是返回一个新的列表。

    1
    2
    3
    import heapq
    arr = [5, 3, 8, 1]
    heapq.heapify(arr)  # 将列表 arr 原地转为堆
    1
    2
    3
    4
    5
    6
    # 复杂度对比
    heapq.heapify(arr)   #O(n) 从最后一个非叶子节点开始，对每个节点做 “下沉” 操作，把它调整为堆

    h = []
    for x in arr:
        heappush(h, x)	#O(nlogn) 每次插入单独调整

    对于自定义对象，需要定义Python 魔术方法: __lt__ 来指定元素之间的比较规则:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    import heapq

    class ListNode:
        def __init__(self, val=0, next=None):
            self.val = val
            self.next = next
            
        def __lt__(self,other):
            return self.val < other.val
        
    # 动态给 ListNode 类增加（或重写）“小于”比较方法:
    ListNode.__lt__ = lambda a,b: a.val < b.val 

    h = [ListNode(1),ListNode(2),ListNode(3)]
    heapq.heapify(h)

  2. heapq.heappush(heap, item)

    将新元素 item 加入堆，并维护堆性质（item放尾部上浮），复杂度：O(log n)

  3. heapq.heappop(heap)

    删除并返回最小值heap[0]，并维护堆性质（最后一个item放堆顶下沉）。复杂度：O(log n)

  4. heapq.heappushpop(heap, item)

    先插入元素再弹出最小值，复杂度：O(log n)

  5. heapq.heapreplace(heap, item)

    先弹出最小值，再插入新值(堆顶一定变化)，复杂度：O(log n)

  6. heapq.nlargest(k, heap) & heapq.nsmallest(k, heap)

    快速找前 k 大/小元素，复杂度：O(n log k)

堆排序 O(nlogn)

1
2
3
4
5
6
7
8
9
10
import heapq
my_list = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
heapq.heapify(my_list)
 
# 依次取出堆中的最小元素
sorted_list = []
while my_list:
    sorted_list.append(heapq.heappop(my_list))
    
print(sorted_list)

find TOP-k small with Max Heap

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import heapq

def topk_small(nums, k):
    if k <= 0:
        return []

    # 1️⃣ 前 K 个元素建大根堆（用负号模拟）
    heap = [-num for num in nums[:k]]
    heapq.heapify(heap)  # O(k)

    # 2️⃣ 遍历剩余元素
    for num in nums[k:]:
        if -num > heap[0]:  # 新元素比堆顶小
            heapq.heapreplace(heap, -num)  # O(log k)

    # 3️⃣ 返回正数
    return [-x for x in heap]

优先队列

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import heapq

pq = [(2, 'task2'), (1, 'task1'), (3, 'task3')] #元组可以按字典序比大小
heapq.heapify(pq)  # O(n) 构建堆

heapq.heappop(pq)  # (1, 'task1')
heapq.heappop(pq)  # (2, 'task2')
heapq.heappop(pq)  # (3, 'task3')

# 也可以自定义优先队列
class Task:
    def __init__(self, name, priority):
        self.name = name
        self.priority = priority

    def __lt__(self, other):
        return self.priority < other.priority  # 用 priority 排序

pq = [Task("task1", 2), Task("task2", 1)]
heapq.heapify(pq)