python堆排序heapq

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标准库》

你可能感兴趣的:(python堆排序heapq)