Python笔记：heapq库简介

最新推荐文章于 2024-06-26 19:29:00 发布

Espresso Macchiato

最新推荐文章于 2024-06-26 19:29:00 发布

阅读量4.2k

点赞数 8

分类专栏： python基础文章标签： python

本文链接：https://blog.csdn.net/codename_cys/article/details/108422407

版权

python基础专栏收录该内容

19 篇文章 6 订阅

订阅专栏

Python笔记：heapq库简介

1. heapq库是什么

heapq库算是一个黑科技，他在原理上并不复杂，事实上就是一个小顶堆结构，即将其转换为一个二叉树结构，则对于每一棵树而言，永远都有叶子节点的值大于根节点的值。

如此，我们只需要要在pop和push时维护好堆结构，就能够保证列表的第一个元素永远是最小的元素。

因此，heapq库的push与pop操作的时间复杂度都是 $O (l o g N)$ ，但是对于需要频繁的插入且取用最小值的情况，heapq库可以大大地提升代码的执行效率。

2. 内置函数

heapq库的内置函数事实上和bisect库一样，也并不多，只有8个，分别为：

heappush(heap, item)
- 向heapq序列中插入元素，时间复杂度 $O (l o g N)$ ；
heappop(heap)
- 从heapq序列中弹出第一个元素（最小值），时间复杂度 $O (l o g N)$ ；
heapify(arr)
- 将一个序列转换为heapq序列，时间复杂度 $O (N l o g N)$ ；
heappushpop(heap, item)
- 先将item元素插入到heapq序列当中，而后弹出最小元素，时间复杂度 $O (l o g N)$ ；
heapreplace(heap, item)
- 先弹出当前的最小元素，而后将item插入到heapq序列当中，时间复杂度 $O (l o g N)$ ；
merge(*iterables, key=None, reverse=False)
- 将多个序列合并为一个heapq序列，时间复杂度 $O (l o g N)$
nlargest(n, iterable, key=None)
- 从一个长序列中取出最大的n个元素，时间复杂度 $O (N l o g n)$ ，等价于sorted(iterable, key=None, reverse=True)[:n]；
nsmallest(n, iterable, key=None)
- 从一个长序列中取出最小的n个元素，，时间复杂度 $O (N l o g n)$ ，等价于sorted(iterable, key=None)[:n]

其中，最为核心的函数只有两个，就是heappush()以及heappop()两个方法，其余的函数本质上来说都能够用这两个函数来改写。且其实现方式也都比较直观，因此，这里就不再赘述了。

我们重点考察一下heappush()和heappop()两个函数的原理和实现。

3. heappop & heappush函数详细考察

如前所述，heapq库中最核心的函数事实上就是元素的插入函数heappush()以及元素的删除函数heappop()，因此，这里，我们就来重点考察一下这两个函数的原理和具体实现。

1. heappush函数

heappush函数包含两个步骤，即：

插入元素；
维护小顶堆结构。

其实现方式也相对较为简单，事实上就是首先将元素插入到序列尾端，然后不断堆的最下方向上调整，直到堆重新满足小顶堆结构。

给出我们自己的代码实现如下：

def heappush(heapq, item):
    n = len(heapq)
    heapq.append(item)
    while n != 0:
        flag = (n-1) // 2
        if heapq[n] < heapq[flag]:
            heapq[flag], heapq[n] = heapq[n], heapq[flag]
            n = flag
        else:
            break
    return

上述代码积即为一种简易的实现，实测得到与heapq中的heappush函数插入结果相同。

2. heappop函数

heappop函数同样可以拆分为两步：

将第一个元素弹出；
维护小顶堆结构

但是，与push操作不同，pop操作后的维护稍微会更加复杂一点，它同样包含几个步骤：

顶部元素的持续补位，即每次将第 $2 k + 1$ 与 $2 k + 2$ 个元素中的较小值（不妨设为第 $2 k + 1$ ）补位到第 $k$ 个元素的位置，而后维护用于补位的第 $2 k + 1$ 个元素；
最后一个元素我们用最后一位的元素进行补位，而后这个补位可能会导致小顶堆结构的破坏，因此，我们需要对这个元素重新像push操作一下维护一下小顶堆结构；
删除最后一个元素；

给出代码实现如下：

def my_heappop(heapq):
    item = heapq[0]
    n = len(heapq)
    flag = 0
    while 2*flag+1 < n:
        if 2*flag+2 >= n or heapq[2*flag+1] < heapq[2*flag+2]:
            heapq[flag] = heapq[2*flag+1]
            flag = 2*flag+1
        else:
            heapq[flag] = heapq[2*flag+2]
            flag = 2*flag+2
    
    # 对最后一个元素进行补位后重新维护一下
    heapq[flag] = heapq[-1]
    while flag != 0:
        if heapq[flag] < heapq[(flag-1) // 2]:
            heapq[flag], heapq[(flag-1) // 2] = heapq[(flag-1) // 2], heapq[flag]
            flag = (flag-1) // 2
        else:
            break
            
    heapq.pop()
    return item