Python堆排序内置模块heapq

简介

该模块提供了堆排序算法的实现。
堆是一颗特殊的完全二叉树。
小根堆：满足任一节点都比其孩子节点小
大根堆：满足任一节点都比其孩子节点大

建堆

方法一: heapify

使用heap.heapify(list)转换列表成为堆结构
import heapq
import random

ls = list(range(10))
random.shuffle(ls)
heapq.heapify(ls)  # 注意，heapq建立的是小根堆
print(ls)

方法二: heappush

使用一个空列表，然后使用heapq.heappush()函数把值加入堆中
import heapq
import random

heap = []
ls = list(range(10))
random.shuffle(ls)
for item in ls:
    heapq.heappush(heap, item)
print(heap)

挨个出数

import heapq
import random

ls = list(range(10))
random.shuffle(ls)
heapq.heapify(ls)  

for i in range(len(ls)):
    print(heapq.heappop(ls))

合并

将多个排序后的序列合并成一个从小到大的序列，返回生成器
import heapq

num1 = [32, 3, 5, 34, 54, 23, 132]
num2 = [23, 2, 12, 656, 324, 23, 54]
num1 = sorted(num1)
num2 = sorted(num2)

res = heapq.merge(num1, num2)
print(res)  # 生成器

获取前N最大或最小的元素

import random

nums = list(range(10))
random.shuffle(nums)

print(heapq.nlargest(n=3, iterable=nums))  # 获取最大
print(heapq.nsmallest(n=3, iterable=nums))  # 获取最小

=====================================================
from pprint import pprint
这两个函数还接受一个key参数，用于dict或其他数据结构类型使用
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': 21.09},
    {'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'])

pprint(cheap)
pprint(expensive)

实现优先队列

heapq的 q 是优先队列的意思

堆的值可以是元组类型，可以实现对带权值的元素进行排序。
from heapq import heappop
from heapq import heappush

h = []
heappush(h, (5, 'write code'))
heappush(h, (7, 'release product'))
heappush(h, (1, 'write spec'))
heappush(h, (3, 'create tests'))
heappush(h, (1, 'bcreate tests'))  # 在heapq中，权值相同, 后添加的先输出
print(heappop(h))

其他方法

其他：
(1) heapreplace: 删除最小元素并且增加一个元素
import heapq
import random

nums = list(range(10))

random.shuffle(nums)

heapq.heapify(nums)

heapq.heapreplace(nums, -1)

print(nums)

(2) heappushpop: 首先判断添加元素值与堆的第一个元素值（堆顶）对比，
                 如果大，则删除第一个元素，然后添加新的元素值，
                 否则不更改堆并且返回添加的值。
print(heapq.heappushpop(nums, 9))  # 返回弹出的最小值
print(heapq.heappushpop(nums, -90))  # 返回添加的值