编程Python

heapq的用法

字数统计: 1.4k字   |   阅读时长: 5分
, ,

摘要: Python 堆操作

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

本文介绍 Python 中与堆相关的组件：heapq 模块。首先介绍建堆、插入、查询极值、删除极值、合并有序序列等操作。之后我们用 heapq 模块解决 2336. 无限集中的最小数字。

Python 的 heapq 模块实现了一个最小堆。算法细节参考文章 二叉堆。关于 Python 中的 PriorityQueue，可以参考 Python标准库-数据结构 中的例子。

建堆

首先构建一个无序列表，然后调用 heapify()。与从空列表开始，一个元素一个元素地调用 heappush() 是一样的。

1
2
data = [9, 8, 7, 6, 5]
heapq.heapify(data)

插入

调用 heappush(data, x)，向堆中插入 x，堆中数据维护在 data 中。插入时保持元素的堆排序性质。

1
2
3
data = []
for x in [9, 8, 7, 6, 5]:
    heapq.heappush(data, x)

删除最小值元素

当堆已经正确组织时，可以用 heappop() 删除最小值元素。

1
smallest = heappop(data)

维持固定大小的堆

如果是删除最小值元素并替换为新值 x，可以用 heapreplace(data, x)。通过这种原地替换元素，可以维持一个固定大小的堆。

1
smallest = heapreplace(data, x)

TopK

heapq 包括两个检查可迭代对象的函数，查找其中包含的最大或最小值的范围。只有当 k 相对较小时才算高效。

1
2
klargest = heapq.nlargest(k, data)
ksmallest = heapq.nsmallest(k, data)

合并有序序列

下面是比较直接的方法，但对于较大数据集会占用大量内存。

1
2
data = [[1, 3, 2], [6, 4, 5]]
merged_list = list(sorted(itertools.chain(*data)))

merge 使用堆一次一个元素地生成新序列。但注意调用 merge 之前，各个序列需要有序。

1
2
3
4
data = [[1, 3, 2], [6, 4, 5]]
for l in data:
    l.sort()
merged_list = list(heapq.merge(*data))

2336. 无限集中的最小数字

现有一个包含所有正整数的集合 [1, 2, 3, 4, 5, …] 。

实现 SmallestInfiniteSet 类：

  • SmallestInfiniteSet() 初始化 SmallestInfiniteSet 对象以包含 所有 正整数。
  • int popSmallest() 移除 并返回该无限集中的最小整数。
  • void addBack(int num) 如果正整数 num 不 存在于无限集中，则将一个 num 添加 到该无限集中。

提示：

1
2
1 <= num <= 1000
最多调用 popSmallest 和 addBack 方法 共计 1000 次

示例：
输入
[“SmallestInfiniteSet”, “addBack”, “popSmallest”, “popSmallest”, “popSmallest”, “addBack”, “popSmallest”, “popSmallest”, “popSmallest”]
[[], [2], [], [], [], [1], [], [], []]
输出
[null, null, 1, 2, 3, null, 1, 4, 5]
解释
SmallestInfiniteSet smallestInfiniteSet = new SmallestInfiniteSet();
smallestInfiniteSet.addBack(2); // 2 已经在集合中，所以不做任何变更。
smallestInfiniteSet.popSmallest(); // 返回 1 ，因为 1 是最小的整数，并将其从集合中移除。
smallestInfiniteSet.popSmallest(); // 返回 2 ，并将其从集合中移除。
smallestInfiniteSet.popSmallest(); // 返回 3 ，并将其从集合中移除。
smallestInfiniteSet.addBack(1); // 将 1 添加到该集合中。
smallestInfiniteSet.popSmallest(); // 返回 1 ，因为 1 在上一步中被添加到集合中，
// 且 1 是最小的整数，并将其从集合中移除。
smallestInfiniteSet.popSmallest(); // 返回 4 ，并将其从集合中移除。
smallestInfiniteSet.popSmallest(); // 返回 5 ，并将其从集合中移除。

算法：堆 + 哈希表

由于每次调用 popSmallest() 需要返回最小的整数。因此可以考虑用最小堆。

由于数据的删除只能通过 popSmallest() 删除最小值，因此当前调用 popSmallest() 如果要删除 $x$，则说明 $[1, n-1]$ 此前都被删除过，并且后来没有插入回来。

当 $x$ 被删后，集合中下一个最小的元素不一定是 $x+1$，因为有可能之前已经删除到了 $x+1$，后来 $x$ 又插入回来了，所以才会使得在 $x+1$ 已经被删的情况下，当前删除值为 $x$。

综上分析，先被删除，后来有插入回来的元素是比较关键的，这部分元素对当前集合的最小值直接造成影响。因此可以用最小堆维护这部分数据。被删除尚未插入回来的元素记录在哈希表中。

此外，用一个变量记录已经删除过的最大值 $x_{max}$，将 $x_{max} + 1$ 也维护到堆中，如果 $x_{max} + 1$ 被删，那么此时堆中一定是空的，此时继续将 $x_{max} + 2$ 插入堆中，这样，堆中始终保持至少一个元素，堆中的最小元素就是调用 popSmallest() 返回的值。

插入 x 的时候，如果 x 在哈希表中，则将 x 从哈希表删除，插入堆中。

代码 (Python)

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

class SmallestInfiniteSet:
    def __init__(self):
        self.heap = []
        self.setting = set()
        heapq.heappush(self.heap, 1)

    def popSmallest(self) -> int:
        smallest = heapq.heappop(self.heap)
        self.setting.add(smallest)
        if len(self.heap) == 0:
            heapq.heappush(self.heap, smallest + 1)
        return smallest

    def addBack(self, num: int) -> None:
        if(num not in self.setting):
            return
        self.setting.remove(num)
        heapq.heappush(self.heap, num)

源码 (维护堆性质的插入)

下面是维护堆性质的插入操作的源码，原理参考文章 二叉堆。

1
2
3
4
def heappush(heap, item):
    """Push item onto heap, maintaining the heap invariant."""
    heap.append(item)
    _siftdown(heap, 0, len(heap)-1)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 'heap' is a heap at all indices >= startpos, except possibly for pos.  pos
# is the index of a leaf with a possibly out-of-order value.  Restore the
# heap invariant.
def _siftdown(heap, startpos, pos):
    newitem = heap[pos]
    # Follow the path to the root, moving parents down until finding a place
    # newitem fits.
    while pos > startpos:
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        if newitem < parent:
            heap[pos] = parent
            pos = parentpos
            continue
        break
    heap[pos] = newitem
Share