heapq模块是python的一个标准库，它实现了一个堆数据结构，堆数据结构是一种二叉树。

  • 什么是堆数据结构？


    • 官网给出的定义是：

    This implementation uses arrays for which
    heap[k] <= heap[2*k+1] and heap[k] <= heap[2*k+2]
    for all k, counting elements from zero.

    我们可以这样理解：
    堆是完全二叉树或者近似二叉树，它的各个节点的键值都有固定对应的的数字，根节点（即root,最上面起始位置）是0，若父节点为heap[k]，则子节点为heap[2*k+1]和heap[2*k+2]。父节点对应的值总是小于或者等于子节点，称为最小堆。对应的，父节点的值总是大于或者等于子节点，称为最大堆。在heapq中，使用的是最小堆。

    正因为堆的这种特殊结构，使得通过heapq模块，可以快速获取一个列表的前N个最大(小)值，即Top N。

  • 堆数据结构的特点？

    • 这里，可能我们不禁要问，python不是内置了sort方法用来排序么？

    现在我们假设一种情景，我们在维护一个列表，并且这个列表在变化，不断有新元素加入，而在任何时候我们可能需要获取里面的Top N，因此我们要求列表始终可以处于已排序状态。

    这时候sort方法就显得不那么好用了，因为每次新加入一个元素，我们都要调用一次sort。数据量小时还是可以的，当数据量较大时，效率就会降低，并且python内置的sort方法本身在实现上也不是那么的高效，复杂度是O(nlog n)。

    python维护了一个堆，使用的存储结构是列表，通过heapq模块来管理、操作这个堆。heapq提供了插入、删除元素的方法，并且保证在插入或删除元素时，所有节点自动调整，保证堆的结构，同时尽量高效，复杂度为O(log n)，在大数据时，效率高于sort排序。

  • 方法描述

    • heapq模块提供了如下几个函数：
    heapq.heappush(heap, item) 把item添加到heap中（heap是一个列表）
    heapq.heappop(heap) 把堆顶元素弹出，返回的就是堆顶
    heapq.heappushpop(heap, item) 先把item加入到堆中，然后再pop，比heappush()再heappop()要快得多
    heapq.heapreplace(heap, item) 先pop，然后再把item加入到堆中，比heappop()再heappush()要快得多
    heapq.heapify(list) 将列表进行堆调整
    heapq.merge(*iterables) 将多个列表合并，并进行堆调整，返回的是合并后的列表的迭代器
    heapq.nlargest(n, iterable, key=None) 返回最大的n个元素（Top-K问题）
    heapq.nsmallest(n, iterable, key=None) 返回最小的n个元素（Top-K问题）

  • heapq使用

    • 获取Top 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]

    使用key关键字参数，进行更为复杂的数据结构处理：

    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'])

    对一个列表进行排序：

    nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
    heapq.heapify(nums)
    print(nums)
    #[-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
  • 总结

    • 堆数据结构的实现有兴趣的可以查看具体的源码实现，堆本身并不是一种有序的结构，但可以通过遍历二叉树的方式得到有序的列表。
    在关于排序和取Top N值时，到底使用什么方法最快，python3 cookbook给出了非常好的建议：


    当要查找的元素个数相对比较小的时候，函数nlargest() 和 nsmallest()。
    仅仅想查找唯一的最小或最大(N=1)的元素的话，那么使用min()和max()函数。
    如果N的大小和集合大小接近的时候，通常先排序这个集合然后再使用切片操作会更快点 (sorted(items)[:N] 或者是 sorted(items)[-N:])。