heapq模块

堆的定义

堆是一种特殊的树形数据结构，每个节点都有一个值，通常我们所说的堆的数据结构指的是二叉树。堆的特点是根节点的值最大（或者最小），而且根节点的两个孩子也能与孩子节点组成子树，亦然称之为堆。 
堆分为两种，大根堆和小根堆是一颗每一个节点的键值都不小于（大于）其孩子节点的键值的树。无论是大根堆还是小根堆（前提是二叉堆）都可以看成是一颗完全二叉树。下面以图的形式直观感受一下: 

heapq模块

在Python中也对堆这种数据结构进行了模块化，我们可以通过调用heapq模块来建立堆这种数据结构，同时heapq模块也提供了相应的方法来对堆做操作。 heapq模块有以下方法：

__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge',
           'nlargest', 'nsmallest', 'heappushpop']

heap = [] #创建了一个空堆 

heappush(heap,item) #往堆中插入一条新的值 

def heappush(heap, item):
    """Push item onto heap, maintaining the heap invariant."""
    heap.append(item)
    _siftdown(heap, 0, len(heap)-1)

heappop(heap) #从堆中弹出最小值 

def heappop(heap):
    """Pop the smallest item off the heap, maintaining the heap invariant."""
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        _siftup(heap, 0)
        return returnitem
    return lastelt

item = heap[0] #查看堆中最小值，不弹出 
heapify(x) #以线性时间讲一个列表转化为堆 

def heapify(x):
    """Transform list into a heap, in-place, in O(len(x)) time."""
    n = len(x)
    # Transform bottom-up.  The largest index there's any point to looking at
    # is the largest with a child index in-range, so must have 2*i + 1 < n,
    # or i < (n-1)/2.  If n is even = 2*j, this is (2*j-1)/2 = j-1/2 so
    # j-1 is the largest, which is n//2 - 1.  If n is odd = 2*j+1, this is
    # (2*j+1-1)/2 = j so j-1 is the largest, and that's again n//2-1.
    for i in reversed(range(n//2)):
        _siftup(x, i)

item = heapreplace(heap,item) #弹出并返回最小值，然后将heapqreplace方法中item的值插入到堆中，堆的整体结构不会发生改变。这里需要考虑到的情况就是如果弹出的值大于item的时候我们可能就需要添加条件来满足function的要求

def heapreplace(heap, item):
    """Pop and return the current smallest value, and add the new item.

   This is more efficient than heappop() followed by heappush(), and can be

    more appropriate when using a fixed-size heap.  Note that the value
    returned may be larger than item!  That constrains reasonable uses of
    this routine unless written as part of a conditional replacement:
        if item > heap[0]:
            item = heapreplace(heap, item)
    """
    return item = heap[0]    # raises appropriate IndexError if heap is empty
    heap[0] = item
    _siftup(heap, 0)
    return returnitem

heappushpop() #将值插入到堆中同时弹出堆中的最小值。 

def heappushpop(heap, item):
    """Fast version of a heappush followed by a heappop."""
    if heap and heap[0] < item:
        item, heap[0] = heap[0], item
        _siftup(heap, 0)
    return item

merge(*iterables) #合并多个堆然后输出

>>> list(heapq.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]
 
 
1
2
 
 

nlargest(n , iterbale, key=None) 
从堆中找出做大的N个数，key的作用和sorted( )方法里面的key类似，用列表元素的某个属性和函数作为关键字。

nsmallest(n, iterable, key=None) #找到堆中最小的N个数

import heapq
nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
print(heapq.nlargest(3, nums)) # Prints [42, 37, 23]
print(heapq.nsmallest(3, nums)) # Prints [-4, 1, 2]

 
 
1
2
3

 
 
1
2
3




两个函数都能接受一个关键字参数，用于更复杂的数据结构中：

 
 


  
  


   
   
portfolio = [
    {'name': 'IBM', 'shares': 100, 'price': 91.1},
    {'name': 'AAPL', 'shares': 50, 'price': 543.22},
    {'name': 'FB', 'shares': 200, 'price': 21.09},
    {'name': 'HPQ', 'shares': 35, 'price': 31.75},
    {'name': 'YHOO', 'shares': 45, 'price': 16.35},
    {'name': 'ACME', 'shares': 75, 'price': 115.65}
]
cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])

  
  

 
 
注意：官方文档中写道使用复合操作方法会比使用单行为操作方法会更快一些。例如使用heappushpop( )方法会比先使用heappush( )再使用heappop( )这两单个方法效率来说会更高。



例子

堆排序如果用到heapq模块就会变得非常好写。首先我们只需要建立一个空堆然后将数导入堆中，再进行弹出操作这样就形成了一个堆排序。
def heapsort(a_iterable):
    H=[]
    for v in a_iterable:
        heapq.heappush(H, v)
    return list(heapq.heappop(L) for i in range(len(H)))
print(heapsort([3, 2, 1, 4, 7, 0, 9, 8, 6]))
[0, 1, 2, 3, 4, 6, 7, 8, 9]

  
  
1
2
3
4
5
6
  
  
1
2
3
4
5
6

在heappush方法的时候heapq就自动给你建立好了一个堆。当然如果对于已有的列表转换成堆也好办，heapq给我们提供了heapify( )方法。
def HeapSort(list):
    heapq.heapify(list)
    heap = []
    while list:
        heap.append(heapq.heappop(list))
    list[:] = heap
    return list
print HeapSort([1,3,5,7,9,2,4,6,8,0])
  
  
1
2
3
4
5
6
7
8
  
  
1
2
3
4
5
6
7
8

当然单个的数列明显不是堆的正常要求，我们可以使用数组的形式来筛选出我们想要的值。这样的话我们能够快速的对字典中我们想要的值做查找，利用刚刚我们讲到的nlargest和nsmallest。
>>> h = []
>>> heappush(h, (5, 'write code'))
>>> heappush(h, (7, 'release product'))
>>> heappush(h, (1, 'write spec'))
>>> heappush(h, (3, 'create tests'))
>>> heappop(h)
(1, 'write spec')
  
  
1
2
3
4
5
6
7
  
  
1
2
3
4
5
6
7


讨论

如果你想在一个集合中查找最小或最大的N个元素，并且N小于集合元素数量，那么这些函数提供了很好的性能。 因为在底层实现里面，首先会先将集合数据进行堆排序后放入一个列表中：

 
 


  
  


   
   
>>> nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
>>> import heapq
>>> heapq.heapify(nums)
>>> nums
[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
>>>


  
  

 
 

堆数据结构最重要的特征是 heap[0] 永远是最小的元素。并且剩余的元素可以很容易的通过调用heapq.heappop() 方法得到，
 该方法会先将第一个元素弹出来，然后用下一个最小的元素来取代被弹出元素(这种操作时间复杂度仅仅是O(log N)，N是堆大小)。 比如，如果想要查找最小的3个元素，你可以这样做：

 
 


  
  


   
   
>>> heapq.heappop(nums)
-4
>>> heapq.heappop(nums)
1
>>> heapq.heappop(nums)
2


  
  

 
 

当要查找的元素个数相对比较小的时候，函数 nlargest() 和 nsmallest() 是很合适的。
 如果你仅仅想查找唯一的最小或最大(N=1)的元素的话，那么使用 min() 和 max() 函数会更快些。
 类似的，如果N的大小和集合大小接近的时候，通常先排序这个集合然后再使用切片操作会更快点 (sorted(items)[:N] 或者是 sorted(items)[-N:] )。
 需要在正确场合使用函数 nlargest() 和nsmallest() 才能发挥它们的优势
 (如果N快接近集合大小了，那么使用排序操作会更好些)。

尽管你没有必要一定使用这里的方法，但是堆数据结构的实现是一个很有趣并且值得你深入学习的东西。 基本上只要是数据结构和算法书籍里面都会有提及到。 heapq 模块的官方文档里面也详细的介绍了堆数据结构底层的实现细节。