heapq

1、如果想实现一个队列,让它以给定的优先级对元素排序,且每次pop操作都返回优先级最高的那个元素,我们可以使用heapq来实现。

下面的类用heapq实现了一个优先级队列:

import heapq
class PriorityQueue:
    def __init__(self):
        self._queue=[]
        self._index=0
    def push(self,item,priority):
        heapq.heappush(self._queue,(-priority,self._index,item))
        self._index+=1
    def pop(self):
        return heapq.heappop(self._queue)[-1]

下面使用这个类:

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

q=PriorityQueue()
q.push(Item('fir'),1)
q.push(Item('sec'),5)
q.push(Item('thi'),4)
q.push(Item('for'),2)
q.push(Item('fiv'),2)

下面看看pop的结果

print(q.pop())
Item('sec')    #pop出优先级最高的
print(q.pop())
Item('thi')    #pop出优先级次高的

对于相同优先级的两个元素,返回顺序和插入顺序相同:

print(q.pop())
Item('for')
print(q.pop())
Item('fiv')

代码分析:

heapq的heappush()和heappop()分别实现将元素从列表中插入和删除,保证列表中第一个元素的优先级最低,heappop()返回优先级最底的元素。上面代码中队列以元组的形式组成(-priority,index,item),priority取负数是为了让队列能够按元素的优先级从高到底顺序排列。index是对于具有相同优先级的元素,按照入队的顺序出队。

2、如果你想在某个集合中找出最大或最小的n个元素,可以使用heapq的nlargest()和nsmallest()。

import heapq
nums=[1,5,4,6,2,7,-1,10,19,-11,22,100]
print(heapq.nlargest(2,nums))
print(heapq.nsmallest(2,nums))
[100, 22]
[-11, -1]

这两个函数还可以接受一个参数key,如下:

portfolio=[
    {'name':'alibaba','shares':100,'price':99.1},
    {'name':'tenx','shares':120,'price':55},
    {'name':'facebook','shares':200,'price':32},
    {'name':'apple','shares':1000,'price':200}
]
cheap=heapq.nsmallest(2,portfolio,key=lambda s:s['price'])
print(cheap)
[{'name': 'facebook', 'shares': 200, 'price': 32}, {'name': 'tenx', 'shares': 120, 'price': 55}]

 

上一篇:Python内置算法与数据结构


下一篇:矩阵快速幂解决斐波那契数问题