一文搞懂heapq

1. Heapq 简介

源码：cpython/Lib/heapq.py at 3.12 · python/cpython (github.com)

heapq模块是python库中实现的小顶堆相关的函数的集合。对外暴露的接口函数有'heappush', 'heappop', 'heapify', 'heapreplace', 'merge', 'nlargest', 'nsmallest', 'heappushpop'。

在Python中，heapq模块提供了对小顶堆（Min Heap）的支持。小顶堆是一种数据结构，其中每个节点的值都小于或等于其子节点的值。这意味着堆的根节点始终是最小的元素。

堆的定义如下，n个关键字序列L[0…n-1] 称为堆，当且仅当该序列满足∶

①L（i）>=L（2i+1）且L（i）>=L（2i+2）或

②L（i）<=L（2i+1）且L（i）<=L（2i+2）（0≤i≤n//2）

可以将该一维数组视为一棵完全二叉树，满足条件①的堆称为大根堆（大顶堆）。大根堆的最大元素存放在根结点。满足条件②的堆称为小根堆(小顶堆)， 小根堆的定义刚好相反， 根结点是最小元素。

2. function method

2.1. heapq.heappush(heap, item)

https://github.com/python/cpython/blob/ff4ad2ea89ed4515340f66657505b6c0deb9a16e/Lib/heapq.py#L132

Push the value item onto the heap, maintaining the heap invariant.

将值项压入堆，保持堆不变性。

2.2. heapq.heappop(heap)

https://github.com/python/cpython/blob/ff4ad2ea89ed4515340f66657505b6c0deb9a16e/Lib/heapq.py#L137

Pop and return the smallest item from the heap, maintaining the heap invariant. If the heap is empty, IndexError is raised. To access the smallest item without popping it, use heap[0].

弹出并返回堆中最小的项，维护堆 不变的。如果堆为空，则引发。要访问 最小的物品不弹出，使用。

2.3. heapq.heappushpop(heap, item)

https://github.com/python/cpython/blob/ff4ad2ea89ed4515340f66657505b6c0deb9a16e/Lib/heapq.py#L163

Push item on the heap, then pop and return the smallest item from the heap. The combined action runs more efficiently than heappush() followed by a separate call to heappop().

Push item到堆上，然后弹出并返回最小的项 堆。联合行动比 接着是另一个电话。

2.4. heapq.heapify(x)

https://github.com/python/cpython/blob/ff4ad2ea89ed4515340f66657505b6c0deb9a16e/Lib/heapq.py#L170

Transform list x into a heap, in-place, in linear time.

在线性时间内将列表x就地转换为堆。

2.5. heapq.heapreplace(heap, item)

https://github.com/python/cpython/blob/ff4ad2ea89ed4515340f66657505b6c0deb9a16e/Lib/heapq.py#L147

Pop and return the smallest item from the heap, and also push the new item. The heap size doesn't change. If the heap is empty, IndexError is raised.

弹出并返回堆中最小的项，同时推入新项。 堆大小不会改变。如果堆为空，则引发。

This one step operation is more efficient than a heappop() followed by heappush() and can be more appropriate when using a fixed-size heap. The pop/push combination always returns an element from the heap and replaces it with item.

这一步操作比后面的操作更有效 并且在使用固定大小的堆时更合适。 pop/push组合总是从堆中返回一个元素并替换 它与item。

The value returned may be larger than the item added. If that isn't desired, consider using heappushpop() instead. Its push/pop combination returns the smaller of the two values, leaving the larger value on the heap.

返回的值可能大于添加的项。如果不是这样的话 渴望，考虑使用替代。其推/流行 组合返回两个值中较小的值，留下较大的值 在堆上。

The module also offers three general purpose functions based on heaps.

该模块还提供了三个基于堆的通用函数。

2.6. heapq.merge(*iterables, key=None, reverse=False)

https://github.com/python/cpython/blob/ff4ad2ea89ed4515340f66657505b6c0deb9a16e/Lib/heapq.py#L316

Merge multiple sorted inputs into a single sorted output (for example, merge timestamped entries from multiple log files). Returns an iterator over the sorted values.

将多个已排序的输入合并为一个已排序的输出(例如，Merge) 来自多个日志文件的带有时间戳的条目)。返回迭代器 遍历排序后的值。

Similar to sorted(itertools.chain(*iterables)) but returns an iterable, does not pull the data into memory all at once, and assumes that each of the input streams is already sorted (smallest to largest).

类似于but返回一个可迭代对象，does 不是一次把所有的数据拉到内存中，并假设每个输入 流已经排序(从最小到最大)。

Has two optional arguments which must be specified as keyword arguments.

有两个可选参数，必须指定为关键字参数。

key specifies a key function of one argument that is used to extract a comparison key from each input element. The default value is None (compare the elements directly).

Key指定一个参数的键函数，用于 从每个输入元素提取一个比较键。默认值为 (直接比较元素)。

reverse is a boolean value. If set to True, then the input elements are merged as if each comparison were reversed. To achieve behavior similar to sorted(itertools.chain(*iterables), reverse=True), all iterables must be sorted from largest to smallest.

Reverse是一个布尔值。如果设置为，则输入元素 被合并，就好像每个比较都是反向的一样。达到相似的行为 To，所有可迭代对象必须 从大到小排序。

Changed in version 3.5: Added the optional key and reverse parameters.

在3.5版更改:添加了可选的key和reverse参数。

2.7. heapq.nlargest(n, iterable, key=None)

https://github.com/python/cpython/blob/ff4ad2ea89ed4515340f66657505b6c0deb9a16e/Lib/heapq.py#L523

Return a list with the n largest elements from the dataset defined by iterable. key, if provided, specifies a function of one argument that is used to extract a comparison key from each element in iterable (for example, key=str.lower). Equivalent to: sorted(iterable, key=key, reverse=True)[:n].

返回一个包含数据集中n个最大元素的列表 可迭代的。如果提供了Key，则指定一个有一个参数的函数 用于从iterable中的每个元素提取比较键(例如， )。相当于:。

2.8. heapq.nsmallest(n, iterable, key=None)

https://github.com/python/cpython/blob/ff4ad2ea89ed4515340f66657505b6c0deb9a16e/Lib/heapq.py#L463

Return a list with the n smallest elements from the dataset defined by iterable. key, if provided, specifies a function of one argument that is used to extract a comparison key from each element in iterable (for example, key=str.lower). Equivalent to: sorted(iterable, key=key)[:n].

返回一个包含数据集中n个最小元素的列表 可迭代的。如果提供了Key，则指定一个有一个参数的函数 用于从iterable中的每个元素提取比较键(例如， )。相当于:。

The latter two functions perform best for smaller values of n. For larger values, it is more efficient to use the sorted() function. Also, when n==1, it is more efficient to use the built-in min() and max() functions. If repeated usage of these functions is required, consider turning the iterable into an actual heap.

后两个函数在较小的n值下表现最好 值时，使用该函数效率更高。另外,当 ，使用内置的和会更有效率 功能。如果需要重复使用这些功能，请考虑转动 将可迭代对象放入实际的堆中。

3. 应用

3.1. 建立大、小根堆

import heapq
a = []   #创建一个空堆
heapq.heappush(a,18)
heapq.heappush(a,1)
heapq.heappush(a,20)
heapq.heappush(a,10)
heapq.heappush(a,5)
heapq.heappush(a,200)
print(a)
>>> [1, 5, 20, 18, 10, 200]


# 建大根堆
a = []
for i in [1, 5, 20, 18, 10, 200]:
    heapq.heappush(a,-i)
print(list(map(lambda x:-x,a)))
>>> [200, 18, 20, 1, 10, 5]

heapq.heapfy()是以线性时间讲一个列表转化为小根堆

a = [1, 5, 20, 18, 10, 200]
heapq.heapify(a)
print(a)

>>> [1, 5, 20, 18, 10, 200]

3.2. 弹出并返回最小的值

import heapq
def heap_sort(arr):
    if not arr:
        return []
    h = []  #建立空堆
    for i in arr:
        heapq.heappush(h,i) #heappush自动建立小根堆
    return [heapq.heappop(h) for i in range(len(h))]  #heappop每次删除并返回列表中最小的值

若是从大到小排列，有两种方法：

1）先建立小根堆，然后每次heappop()，此时得到从小大的排列，再reverse

2）利用相反数建立大根堆，然后heappop(-元素)。即push(-元素)，pop(-元素)

3.3. 先存再取

heapq.heappushpop()是heappush和haeppop的结合，同时完成两者的功能，先进行heappush()，再进行heappop()

>>>h =  [1, 2, 9, 5]
>>> heappop(h)
1
>>> heappushpop(h,4)            #增加4同时删除最小值2并返回该最小值，与下列操作等同：
2                              
>>> h
[4, 5, 9]

3.4. 先取再存

heapq.heapreplace()与heapq.heappushpop()相反，先进行heappop()，再进行heappush()

堆的大小不变。 如果堆为空则引发 IndexError。这个单步骤操作比依次执行heappop() + heappush() 更高效，并且在使用固定大小的堆时更为适宜。 pop/push 组合总是会从堆中返回一个元素并将其替换为 item。返回的值可能会比添加的 item 更大。 如果不希望如此，可考虑改用 heappushpop()。 它的 push/pop 组合会返回两个值中较小的一个，将较大的值留在堆中。

>>> a=[]
>>> heapreplace(a,3)            #如果list空，则报错
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
IndexError: index out of range
    >>> heappush(a,3)
    >>> a
    [3]
    >>> heapreplace(a,2)            #先执行删除（heappop(a)->3),再执行加入（heappush(a,2))
    3
    >>> a
    [2]
    >>> heappush(a,5)  
    >>> heappush(a,9)
    >>> heappush(a,4)
    >>> a
    [2, 4, 9, 5]
    >>> heapreplace(a,6)            #先从堆a中找出最小值并返回，然后加入6
    2
    >>> a
    [4, 5, 9, 6]
    >>> heapreplace(a,1)            #1是后来加入的，在1加入之前，a中的最小值是4
    4
    >>> a
    [1, 5, 9, 6]

3.5. 合并

heapq.merge()合并多个堆然后输出

输入的list无序，merge后无序，若输入的list有序，merge后也有序

list(heapq.merge([1578,2,5,3],[1564,554,25458],])) #无序

>>> [1564,554,1578,2,5,3,25458]

list(merge([1,3,5,7],[0,2,4,8],[5,10,15,20],[],[25])) #有序

>>> [0,1,2,3,4,5,5,7,8,10,15,20,25]

3.6. 获取前 top 值

获取列表中最大、最小的几个值，key的作用和sorted( )方法里面的key类似

>>>a = [0, 1, 2, 3, 4, 5, 5, 7, 8, 10, 15, 20, 25]
>>>heapq.nlargest(5,a)
[25, 20, 15, 10, 8]

>>>b = [('a',1),('b',2),('c',3),('d',4),('e',5)]
>>>heapq.nlargest(1,b,key=lambda x:x[1])
[('e', 5)]

3.7. 复杂度

1）heapq.heapify(x): O(n)

2）heapq.heappush(heap, item): O(logn)

3）heapq.heappop(heap): O(logn)

即插入或删除元素时，所有节点自动调整，保证堆的结构的复杂度为O(log n)

4）heapq.nlargest(k,iterable)和nsmallest(k,iterable)：O(n * log(t))

查找的注意点：

在关于排序和取Top N值时，到底使用什么方法最快，python3 cookbook给出了非常好的建议：

1）当要查找的元素个数相对比较小的时候，函数nlargest() 和 nsmallest()。

2）仅仅想查找唯一的最小或最大(N=1)的元素的话，那么使用min()和max()函数。

3）如果N的大小和集合大小接近的时候，通常先排序这个集合然后再使用切片操作会更快点 (sorted(items)[:N] 或者是 sorted(items)[-N:])。

4. 算法应用

4.1. 实现优先队列

# priority 优先级

class PriorityQueue:
    def __init__(self):
        self._queue = []
        self._index = 0
    def push(self, item, priority):
        # heappush 在队列 _queue 上插入第一个元素
        heapq.heappush(self._queue, (-priority, self._index, item))
        self._index += 1
    def pop(self):
        # heappop 在队列 _queue 上删除第一个元素
        return heapq.heappop(self._queue)[-1]

class Item:
    def __init__(self, name):
        self.name = name
    def __repr__(self):
        return 'Item({})'.format(self.name)

调用push（）方法，实现将列表转化为堆数据。

插入的是元组，元组大小比较是从第一个元素开始，第一个相同，再对比第二个元素，我们这里采用的方案是如果优先级相同，那么就根据第二个元素，谁先插入堆中，谁的index就小，那么它的值就小。

heapq.heappop() 方法得到，该方法会先将第一个元素弹出来，然后用下一个最小的元素来取代被弹出元素。

4.2. 优先队列测试

q = PriorityQueue()
q.push(Item('foo'), 1)
q.push(Item('bar'), 5)
q.push(Item('spam'), 4)
q.push(Item('grok'), 1)

print(q.pop())
print(q.pop())
print(q.pop())

>>> Item('bar')
>>> Item('spam')
>>> Item('foo')

4.3. 实现任务优先级

为了实现这样一个队列，需要跟踪每个元素的位置，并且在更新优先级时调整元素的位置。

import heapq

class PriorityQueue:
    REMOVED = '<removed-task>'  # placeholder for a removed task
    counter = 0  # unique sequence count

    def __init__(self):
        self.pq = []                         # list of entries arranged in a heap
        self.entry_finder = {}               # mapping of tasks to entries
        self.counter = itertools.count()     # unique sequence count

    def add_task(self, task, priority=0):
        'Add a new task or update the priority of an existing task'
        if task in self.entry_finder:
            self.remove_task(task)
        count = next(self.counter)
        entry = [priority, count, task]
        self.entry_finder[task] = entry
        heapq.heappush(self.pq, entry)

    def remove_task(self, task):
        'Mark an existing task as REMOVED.  Raise KeyError if not found.'
        entry = self.entry_finder.pop(task)
        entry[-1] = PriorityQueue.REMOVED

    def pop_task(self):
        'Remove and return the lowest priority task. Raise KeyError if empty.'
        while self.pq:
            priority, count, task = heapq.heappop(self.pq)
            if task is not self.REMOVED:
                del self.entry_finder[task]
                return task
        raise KeyError('pop from an empty priority queue')

    def change_priority(self, task, priority):
        'Change the priority of an existing task'
        self.remove_task(task)
        self.add_task(task, priority)

这个类实现了一个可以改变任务优先级的优先队列。当任务的优先级改变时，我们先将其标记为删除，然后重新添加到队列中。这样可以保持堆的性质不变。

4.4. 实现延迟任务

延迟队列（Delayed Queue）是一种特殊的队列类型，它要求队列中的元素只有在达到了指定的延迟时间后才能被消费。在Python中，我们可以利用heapq模块来构建这样的队列。每个元素不仅包含实际的任务信息，还包括该任务何时可以被处理的时间戳。

import heapq
import time

class DelayedQueue:
    def __init__(self):
        self.queue = []
        self.index = 0

    def insert_with_delay(self, item, delay):
        """
        插入一个项目，并设置其延迟时间（以秒为单位）。
        """
        deadline = time.time() + delay
        entry = [deadline, self.index, item]
        heapq.heappush(self.queue, entry)
        self.index += 1

    def get_ready(self):
        """
        获取所有已准备好处理的任务（即当前时间已经晚于它们的截止时间）。
        返回一个列表，包含所有已准备好处理的任务。
        """
        now = time.time()
        ready = []
        while self.queue and self.queue[0][0] <= now:
            _, _, item = heapq.heappop(self.queue)
            ready.append(item)
        return ready

# 使用示例
dq = DelayedQueue()
dq.insert_with_delay("Task A", 5)  # 设置5秒后的延迟
dq.insert_with_delay("Task B", 10)  # 设置10秒后的延迟
dq.insert_with_delay("Task C", 0)  # 立即可用

# 假设现在时间已经过了5秒
time.sleep(5)

# 获取所有准备好的任务
print(dq.get_ready())  # 应该输出 ["Task C", "Task A"]

在这个实现中，insert_with_delay 方法将任务和延迟时间加入到队列中。每个任务都被包装成一个三元组 [deadline, index, item]，其中 deadline 是任务可以被处理的时间戳，index 是一个递增的整数，用来打破相同时间戳任务之间的顺序，确保先进先出的原则；item 是实际要处理的任务。

get_ready 方法会检查当前时间是否已经超过了队列中最紧迫任务的截止时间。如果是，则从队列中弹出所有可处理的任务并返回它们。

注意：在实际的应用中，如果延迟队列需要长时间运行并且可能包含大量条目，那么在设计时还需要考虑到内存管理的问题。此外，如果需要支持任务的取消或修改延迟时间等功能，那么实现就会更加复杂。

4.5. 实现双端优先队列

双端优先队列（Two-ended Priority Queue 或者 Min-Max Heap）是一种特殊的数据结构，它允许我们在两端进行插入和删除操作，即可以在队列的一端插入最小值，在另一端插入最大值，并且可以从两端分别取出最小值和最大值。这种数据结构结合了优先队列和双端队列（deque）的特点。

在Python中，我们可以使用heapq模块来实现一个简单的双端优先队列。由于heapq模块主要支持最小堆，我们需要额外的逻辑来支持最大堆的功能。一种方法是通过存储每个元素的负数来实现最大堆的行为。

import heapq

class MinMaxHeap:
    def __init__(self):
        self.min_heap = []  # 用于最小堆
        self.max_heap = []  # 用于最大堆

    def push_min(self, value):
        """ 向最小堆中添加一个元素 """
        heapq.heappush(self.min_heap, value)

    def push_max(self, value):
        """ 向最大堆中添加一个元素 """
        heapq.heappush(self.max_heap, -value)

    def pop_min(self):
        """ 从最小堆中弹出并返回最小元素 """
        if self.min_heap:
            return heapq.heappop(self.min_heap)
        else:
            raise IndexError("pop from an empty min heap")

    def pop_max(self):
        """ 从最大堆中弹出并返回最大元素 """
        if self.max_heap:
            return -heapq.heappop(self.max_heap)
        else:
            raise IndexError("pop from an empty max heap")

    def peek_min(self):
        """ 查看最小堆的最小元素，不弹出 """
        if self.min_heap:
            return self.min_heap[0]
        else:
            raise IndexError("peek from an empty min heap")

    def peek_max(self):
        """ 查看最大堆的最大元素，不弹出 """
        if self.max_heap:
            return -self.max_heap[0]
        else:
            raise IndexError("peek from an empty max heap")

# 示例
m = MinMaxHeap()
m.push_min(3)
m.push_min(1)
m.push_max(7)
m.push_max(5)

print(m.peek_min())  # 输出: 1
print(m.peek_max())  # 输出: 7
print(m.pop_min())   # 输出: 1
print(m.pop_max())   # 输出: 7

在这个实现中，push_min 和 pop_min 分别用于向最小堆中添加元素和弹出最小元素，而 push_max 和 pop_max 则用于向最大堆中添加元素和弹出最大元素。为了实现最大堆，我们在插入元素时取反，这样最小堆的操作就变成了寻找最大值的操作。

需要注意的是，这个简单的实现没有处理元素在两个堆之间平衡的问题，也没有处理元素重复的情况。在某些应用场景下，可能需要更复杂的逻辑来维护堆之间的平衡，以及处理重复元素的插入和删除。

如果你需要一个完全平衡的双端优先队列，可能需要更复杂的实现，例如使用专门的数据结构如Min-Max Heap或者Treap等，这些数据结构可以更高效地支持两端操作。