Python中heapq与优先队列【详细】

本文始发于我的公众号：TechFlow, 原创不易，求个关注

今天的文章来介绍Python当中一个蛮有用的库——heapq。

heapq的全写是heap queue，是堆队列的意思。这里的堆和队列都是数据结构，在后序的文章当中咱们会详细介绍，今天只介绍heapq的用法，若是不了解heap和queue原理的同窗能够忽略，咱们并不会深刻太多，会在以后的文章里详细阐述。

在介绍用法以前，咱们须要先知道优先队列的定义。队列你们应该都不陌生，也是很是基础简单的数据结构。咱们能够想象成队列里的全部元素排成一排，新的元素只能从队尾加入队列，元素要出队列只能经过队首，不能中途从队列当中退出。而优先队列呢，是给队列当中的元素每个都设置了优先级，使得队伍当中的元素会自动按照优先级排序，优先级高的排在前面。

也就是说Python当中的heapq就是一个维护优先队列的library，咱们经过调用它能够轻松实现优先队列的功能。

最大或最小的K个元素

咱们来看一个实际的问题，假设咱们当下有N个杂乱无章的元素，可是咱们只关心其中最大的K个或者是最小的K个元素。咱们想从整个数组当中将这部分抽取出来，应该怎么办呢？

这个问题在实际当中很是常见，随便就能够举出例子来。好比用户输入了搜索词，咱们根据用户的搜索词找到了大量的内容。咱们想要根据算法筛选出用户最有可能点击的文原本，机器学习的模型能够给每个文本一个预测的分数。以后，咱们就须要选出分数最大的K个结果。这种相似的场景还有不少，利用heapq库里的nlargest和nsmallest接口能够很是方便地作到这点。

咱们一块儿来看一个例子：

import heapq

nums = [14, 20, 5, 28, 1, 21, 16, 22, 17, 28]
heapq.nlargest(3, nums)
# [28, 28, 22]
heapq.nsmallest(3, nums)
# [1, 5, 14]
复制代码

heapq的nlargest和nsmallest接受两个参数，第一个参数是K，也就是返回的元素的数量，第二个参数是传入的数组，heapq返回的正是传入的数组当中的前K大或者是前K小。

这里有一个问题，若是咱们数组当中的元素是一个对象呢？应该怎么办？

其实也很简单，有了解过Python自定义关键词排序的同窗应该知道，和排序同样，咱们能够经过匿名函数实现。

匿名函数

咱们都知道，在Python当中经过def能够定义一个函数。经过def定义的函数都有函数名，因此称为有名函数。除了有名函数以外，Python还支持匿名函数。顾名思义，就是没有函数名的函数。也就是说它其余方面都和普通函数同样，只不过没有名字而已。

初学者可能会纳闷，函数没有名字应该怎么调用呢？

会有这个疑惑很正常，这是由于习惯了面向过程的编程，对面向对象理解不够深刻致使的。在许多高级语言当中，一切皆对象，一个类，一个函数，一个int都是对象。既然函数也是对象，那么函数天然也能够用来传递，不只能够用来传递，还能够用来返回。这是函数式编程的概念了，咱们这里很少作深刻。

固然，普通函数也同样能够传递，起到的效果同样。只不过在编程当中，有些函数咱们只会使用一次，不必再单独定义一个函数，使用匿名函数会很是方便。

举个例子，比方说我有一个这样的函数：

def operate(x, func):
  return func(x)
复制代码

这个operate函数它接受两个参数，第一个参数是变量x，第二个参数是一个函数。它会在函数内部调用func，返回func调用的结果。我如今要作这样一件事情，我但愿根据x这个整数对4取余的余数来判断应该用什么样的func。若是对4的余数为0，我但愿求一次方，若是余数是2，我但愿求平方，以此类推。若是按照正常的方法，咱们须要实现4个方法，而后依次传递。

这固然是能够的，不过很是麻烦，若是使用匿名函数，就能够大大简化代码量：

def get_result(x):
  if x % 4 == 0:
    return operate(x, lambda x: x)
  elif x % 4 == 1:
    return operate(x, lambda x: x ** 2)
  elif x % 4 == 2:
    return operate(x, lambda x: x ** 3)
  else:
    return operate(x, lambda x: x ** 4)
复制代码

在上面的代码当中，咱们经过lambda关键字定义了匿名函数，避免了定义四种函数用来传递的状况。固然，这个问题还有更简单的写法，能够只用一个函数解决。

咱们来看lambda定义匿名函数的语法，首先是lambda关键字，表示咱们当下定义的是一个匿名函数。以后跟的是这个匿名函数的参数，咱们只用到一个变量x，因此只须要写一个x。若是咱们须要用到多个参数，经过逗号分隔，固然也能够不用参数。写完参数以后，咱们用冒号分开，冒号后面写的是返回的结果。

咱们也能够把匿名函数赋值给一个变量，以后咱们就能够和调用普通函数同样来调用了：

square = lambda x: x ** 2

print(square(3))
print(operate(3, square))
复制代码

自定义排序

回到以前的内容，若是咱们想要heapq排序的是一个对象。那么heapq并不知道应该依据对象当中的哪一个参数来做为排序的衡量标准，因此这个时候，须要咱们本身定义一个获取关键字的函数，传递给heapq，这样才能够完成排序。

好比说，咱们如今有一批电脑，咱们但愿heapq可以根据电脑的价格排序：

laptops = [
    {'name': 'ThinkPad', 'amount': 100, 'price': 91.1},
    {'name': 'Mac', 'amount': 50, 'price': 543.22},
    {'name': 'Surface', 'amount': 200, 'price': 21.09},
    {'name': 'Alienware', 'amount': 35, 'price': 31.75},
    {'name': 'Lenovo', 'amount': 45, 'price': 16.35},
    {'name': 'Huawei', 'amount': 75, 'price': 115.65}
]

cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])
复制代码

在调用nlargest和nsmallest的时候，咱们额外传递了一个参数key，咱们传入的是一个匿名函数，它返回的结果是这个对象的price，也就是说咱们但愿heapq根据对象的price来进行排序。

优先队列

heapq除了能够返回最大最小的K个数以外，还实现了优先队列的接口。咱们能够直接调用heapq.heapify方法，输入一个数组，返回的结果是根据这个数组生成的堆（等价于优先队列）。

固然咱们也能够从零开始，直接经过调用heapq的push和pop来维护这个堆。接下来，咱们就经过heapq来本身动手实现一个优先队列，代码很是的简单，我想你们应该能够瞬间学会。

首先是实现优先队列的部分：

import heapq

class PriorityQueue:
  
  def __init__(self):
    self._queue = []
    self._index =0
    
  def push(self, item, priority):
    # 传入两个参数，一个是存放元素的数组，另外一个是要存储的元素，这里是一个元组。
    # 因为heap内部默认有小到大排，因此对priority取负数
    heapq.heappush(self._queue, (-priority, self._index, item))
    self._index += 1
  
  def pop(self):
    return heapq.heappop(self._queue)[-1]
复制代码

其次咱们来实际看一下运用的状况：

q = PriorityQueue()

q.push('lenovo', 1)
q.push('Mac', 5)
q.push('ThinkPad', 2)
q.push('Surface', 3)

q.pop()
# Mac
q.pop()
# Surface
复制代码

到这里，关于heapq的应用方面就算是介绍完了，可是尚未真正的结束。

咱们须要分析一下heapq当中操做的复杂度，关于堆的部分咱们暂时跳过，咱们先来看nlargest和nsmallest。我在github当中找到了这个库的源码，在方法的注释上，做者写下了这个方法的复杂度，和排序以后取前K个开销五五开：

def nlargest(n, iterable, key=None):
    """Find the n largest elements in a dataset. Equivalent to: sorted(iterable, key=key, reverse=True)[:n] """
复制代码

咱们都知道排序的复杂度的指望是，若是你了解堆的话，会知道堆一次插入元素的复杂度是。若是咱们限定堆的长度是K，咱们插入n次以后也只能保留K个元素。每次插入的复杂度是，一共插入n次，因此总体的复杂度是。

若是K小一些，可能开销会比排序稍小，可是程度有限。那么有没有什么办法能够不用排序而且尽量快地筛选出前K大或者是前K小的元素呢？

我这里先卖个关子，咱们以后的文章当中再来说解。

今天的文章就到这里，若是以为有所收获，请顺手点个关注吧，你的举手之劳对我很重要。

参考资料

Python CookBook Version3

维基百科