Python笔记：heapq库简介

• Python笔记：heapq库简介
  • 1. heapq库是什么
  • 2. 内置函数
  • 3. heappop & heappush函数详细考察
    • 1. heappush函数
    • 2. heappop函数
  • 4. 参考链接

1. heapq库是什么

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

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

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

2. 内置函数

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

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

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

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

3. heappop & heappush函数详细考察

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

1. heappush函数

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

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

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

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

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函数同样可以拆分为两步：

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

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

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

给出代码实现如下：

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

经测试，上述代码实现与heapq库中的heappop函数具有相同的实现结果表现。

4. 参考链接

1. https://docs.python.org/3/library/heapq.html