heapq制作优先级队列

heapq制作优先级队列

1.概述

python里的队列都是先进先出队列，例如list创建的队列。他们会按照接收元素的顺序保存这些元素。
但有时候，我们想根据元素的重要程度来排序，在这种情况下，应该使用优先级队列。
使用堆可以实现优先级队列，那么列表实现的队列和堆实现的列表有什么区别那？

• 列表实现的队列默认存储元素的顺序是添加元素的顺序，并且列表底层使用数组结构存储数据效率不高。
• 堆实现队列不是按照添加元素的顺序存储，而是按照指定的属性排序存储，它底层使用二叉树存储数据，效率非常高。

2.使用普通队列存在的问题

2.1.开发一个图书管理功能

下面我们通过实现一个图书馆借书功能演示使用普通队列的问题。
假如我们要给图书馆写个程序，管理图书的出借和归还的工作。有的人能及时还书，有的人不按时还书，所以我们要提醒他们。下面是结束和还书的代码。

''' 创建图书对象'''
# Example 1
class Book:
    def __init__(self, title, due_date):
        '''

        :param title: 书名
        :param due_date: 归还日期
        '''
        self.title = title
        self.due_date = due_date


# Example 2
def add_book(queue, book):
    queue.append(book)
    # 根据归还日期排序，最早归还日期在列表尾部
    queue.sort(key=lambda x: x.due_date, reverse=True)

# 添加借书和归还日期
queue = []
add_book(queue, Book('Don Quixote', '2019-06-07'))
add_book(queue, Book('Frankenstein', '2019-06-05'))
add_book(queue, Book('Les Misérables', '2019-06-08'))
add_book(queue, Book('War and Peace', '2019-06-03'))

# Example 3
class NoOverdueBooks(Exception):
    pass

''' 返回逾期没有还书的图书'''
def next_overdue_book(queue, now):
    if queue:
        book = queue[-1]
        # 与当前日期比较，按照归还日期从远到近的顺序返回没有归还的书
        if book.due_date < now:
            queue.pop()
            return book
    # 如果没有预期没有归还的则返回消息
    raise NoOverdueBooks


# Example 4
now = '2019-06-10'

# 多次调用输出预期未换的书名
found = next_overdue_book(queue, now)
print(found.title)

found = next_overdue_book(queue, now)
print(found.title)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

运行上面的代码，根据逾期时间的远近关系返回了没有归还的图书

War and Peace
Frankenstein
1
2

2.2.测试图书借书和归还图书性能

1.测试借书性能

借书的过程是不断向列表中添加数据，因此我们来测试下向列表添加数据性能花费的时间就是借书消耗时间。

from timeit import timeit, repeat
import random
import timeit

def print_results(count, tests):
    avg_iteration = sum(tests) / len(tests)
    print(f'Count {count:>5,} takes {avg_iteration:.6f}s')
    return count, avg_iteration

def print_delta(before, after):
    before_count, before_time = before
    after_count, after_time = after
    growth = 1 + (after_count - before_count) / before_count
    slowdown = 1 + (after_time - before_time) / before_time
    print(f'{growth:>4.1f}x data size, {slowdown:>4.1f}x time')

def list_overdue_benchmark(count):
    def prepare():
        to_add = list(range(count))
        random.shuffle(to_add)
        return [], to_add

    def run(queue, to_add):
        for i in to_add:
            # 添加图书，并将列表中的顺序反转
            queue.append(i)
            queue.sort(reverse=True)

    tests = repeat(
        # prepare函数返回list和count
        setup='queue, to_add = prepare()',
        stmt=f'run(queue, to_add)',
        globals=locals(),
        repeat=100,
        number=1)

    return print_results(count, tests)

# Example 8
baseline = list_overdue_benchmark(500)
for count in (1_000, 1_500, 2_000):
    print()
    comparison = list_overdue_benchmark(count)
    print_delta(baseline, comparison)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

运行上面的代码，从花费的时间可以看出它和队列长度呈线性增长。

Count   500 takes 0.001161s

Count 1,000 takes 0.003619s
 2.0x data size,  3.1x time

Count 1,500 takes 0.007473s
 3.0x data size,  6.4x time

Count 2,000 takes 0.012954s
 4.0x data size, 11.2x time
1
2
3
4
5
6
7
8
9
10

2.测试还书性能

还书是从列表中删除数据过程，首先要在列表中找到待待换的图书，然后在列表中删除。这时，位于这条记录之后的数据都要向前移动一个位置，这样算下来，执行还书的时间也呈线性增长。

def list_return_benchmark(count):
    def prepare():
        queue = list(range(count))
        random.shuffle(queue)

        to_return = list(range(count))
        random.shuffle(to_return)

        return queue, to_return

    def run(queue, to_return):
        for i in to_return:
            queue.remove(i)

    tests = repeat(
        setup='queue, to_return = prepare()',
        stmt=f'run(queue, to_return)',
        globals=locals(),
        repeat=100,
        number=1)

    return print_results(count, tests)


# Example 10
baseline = list_return_benchmark(500)
for count in (1_000, 1_500, 2_000):
    print()
    comparison = list_return_benchmark(count)
    print_delta(baseline, comparison)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

运行上面的代码，查看还书消耗的时间呈线性增长

Count   500 takes 0.000919s

Count 1,000 takes 0.003649s
 2.0x data size,  4.0x time

Count 1,500 takes 0.008248s
 3.0x data size,  9.0x time

Count 2,000 takes 0.014748s
 4.0x data size, 16.0x time
1
2
3
4
5
6
7
8
9
10

3.使用优先级队列

使用list做队列保存很少的图书还可以，但是要保存的图书很多时候就会变得运行非常的慢，好在python内置了heapq模块可以解决这个问题，因为它能够高效的实现优先队列。
模块里heap指的是堆，他是一种数据结构，可以维护列表中的元素，并且只需要对数据级别的时间就可以添加或删除其中最小的元素。这种算法复杂度要低于线性算法，所以效率比线性算法高。

3.1.使用heapq优化图书管理功能

1.添加图书到堆数据结构

用heapq模块重新实现add_book函数，这次的队列(queue参数)本身还是个普通列表，但我们不能像原来那样使用list.append方法添加新元素，而是使用heappush函数添加元素。而且也不需要在队列上调用list.sort方法排序了。

如果直接使用原来的Book类，添加数据会报错，这是因为heapq模块规定，添加到优先队列里面的元素必须是可比较的，并且要具备自然顺序，而目前的Book类还不满足这一要求。

python内置的functools模块里有个名为total_ordering的类修饰器，只需要用它修饰Book，并把描述小于关系在魔法函数lt中定义出来。

这样实现之后，就可以通过add_bool函数正常向优先级队列中添加记录了。

from heapq import heappush

def add_book(queue, book):
    heappush(queue, book)

import functools

@functools.total_ordering
class Book:
    def __init__(self, title, due_date):
        self.title = title
        self.due_date = due_date
	# 判断概述的归还日期是不是比另一本书早
    def __lt__(self, other):
        return self.due_date < other.due_date


# Example 14
queue = []
add_book(queue, Book('Pride and Prejudice', '2019-06-01'))
add_book(queue, Book('The Time Machine', '2019-05-30'))
add_book(queue, Book('Crime and Punishment', '2019-06-06'))
add_book(queue, Book('Wuthering Heights', '2019-06-12'))
print([b.title for b in queue])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

运行上面的代码，图书添加到列表了，并且它的顺序是按照还书日期由远到近排序。

['The Time Machine', 'Pride and Prejudice', 'Crime and Punishment', 'Wuthering Heights']
1

2.归还图书从堆中删除元素

使用heappop删除堆中的元素，完成归还图书功能。

from heapq import heappop

# 判断是否有未归还的图书
def next_overdue_book(queue, now):
    if queue:
        book = queue[0]           # Most overdue first
        if book.due_date < now:
            # 使用heappop弹出元素
            heappop(queue)        # Remove the overdue book
            return book

    raise NoOverdueBooks

'''找出预期还没有归还的书,知道全部图书归还'''

# Example 18
now = '2019-06-02'

book = next_overdue_book(queue, now)
print(book.title)

book = next_overdue_book(queue, now)
print(book.title)

try:
    next_overdue_book(queue, now)
except NoOverdueBooks:
    pass          # Expected
else:
    assert False  # Doesn't happen
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

运行上面的代码，输出了预期未归还的图书

The Time Machine
Pride and Prejudice
1
2

3.2.测试heapq队列添加和移除元素性能

def heap_overdue_benchmark(count):
    def prepare():
        to_add = list(range(count))
        random.shuffle(to_add)
        return [], to_add

    def run(queue, to_add):
        for i in to_add:
            heappush(queue, i)
        while queue:
            heappop(queue)

    tests = repeat(
        setup='queue, to_add = prepare()',
        stmt=f'run(queue, to_add)',
        globals=locals(),
        repeat=100,
        number=1)

    return print_results(count, tests)


# Example 20
baseline = heap_overdue_benchmark(500)
for count in (1_000, 1_500, 2_000):
    print()
    comparison = heap_overdue_benchmark(count)
    print_delta(baseline, comparison)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

运行上面的代码，从结果中可以看出消耗的时间比普通的列表快的多

Count   500 takes 0.000172s

Count 1,000 takes 0.000345s
 2.0x data size,  2.0x time

Count 1,500 takes 0.000522s
 3.0x data size,  3.0x time

Count 2,000 takes 0.000946s
 4.0x data size,  5.5x time
1
2
3
4
5
6
7
8
9
10