heapq模块实现了一个适用于Python列表的最小堆排序算法。
堆是一种树形数据结构,其中子节点与父节点之间是一种有序关系。最大堆中父节点大于或等于两个子节点,最小堆父节点小于或等于两个子节点。Python的heapq模块实现了一个最小堆。
创建堆
创建堆有两种方式,heappush()和heapify()。
import heapq
data = [1,5,3,2,8,5]
heap = []
for n in data:
heapq.heappush(heap, n)
如果数据已经存在于内存中,使用heapify()原地重新组织列表中的元素会更加高效。
import heapq
data = [1,5,3,2,8,5]
heapq.heapify(data)
上述两种方法得到的堆是一样的。
还有一种merge()方法,可以从多个可迭代对象中创建一个堆,相当于heapq.heapify(itertools.chain(*iterables))
访问堆内容
堆创建好之后,可以使用heappop()删除有最小值的元素。
import heapq
data = [1,5,3,2,8,5]
heapq.heapify(data)
print heapq.heappop(data) # print 1
如果希望在一个操作中删除现有最小元素并替换成新值,可以使用heapreplace()。
import heapq
data = [1,5,3,2,8,5]
heapq.heapify(data)
heapq.heapreplace(data,10)
堆的数据极值
heapq还包括两个检查可迭代对象的函数,查找其中包含的最大值与最小值的范围。
import heapq
data = [1,5,3,2,8,5]
print heapq.nlargest(3, data)
print heapq.nsmallest(3, data)
输出结果如下:
[8, 5, 5]
[1, 2, 3]
只有在n值比较小的情况下nlargest()和nsmallest()才比较高效。
这个两个函数还接受一个关键字参数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'])
上面代码对每个元素进行比较时, 会以price的值进行比较。
实战
实现堆排序:
>>> from heapq import *
>>> def heapsort(iterable):
... 'Equivalent to sorted(iterable)'
... h = []
... for value in iterable:
... heappush(h, value)
... return [heappop(h) for i in range(len(h))]
...
>>> heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
由于堆中的元素可以为元组,所以可以对带权值的元素进行排序。
>>> from heapq import *
>>> 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')
参考
《python参考手册》
《python标准库》