堆数据结构(heapq)简单应用

## 堆数据结构(heapq)简单应用

# 堆数据结构 heapq
# 常用方法：nlargest(),nsmallest(),heapify(),heappop()
# 如果需要的个数较小，使用nlargest或者nsmallest比较好
# 如果需要的个数已经接近了序列长度，使用sorted()[:N]获取前N个数据比较好
# 如果只需要唯一一个最大或最小值，则直接使用max()或者min()

import heapq


nums = [1,3,5,67,7,34,6,8,5,-12,-45,-234,232]
print(heapq.nlargest(3, nums))
# [232, 67, 34]
print(heapq.nsmallest(3, nums))
# [-234, -45, -12]

words = ['happy', 'sad', 'fun', 'sweet', 'blue']
print(heapq.nlargest(3, words))
# ['sweet', 'sad', 'happy']
print(heapq.nsmallest(3, words))
# ['blue', 'fun', 'happy']

print(heapq.nsmallest(3, 'qazwsxedcvbnm'))
# ['a', 'b', 'c']

t = (1,2,3,4)
print(heapq.nlargest(2, t))
#[4, 3]

students = [
    {"name": "Stanley", "score": 94},
    {"name": "Lily", "score": 98},
    {"name": "Bob", "score": 87},
    {"name": "Well", "score": 85}
]

print(heapq.nlargest(len(students), students, key=lambda s: s["score"])) # 需要个数为序列长度，不好
"""
[{'name': 'Lily', 'score': 98},
{'name': 'Stanley', 'score': 94},
{'name': 'Bob', 'score': 87},
{'name': 'Well', 'score': 85}]
"""

print(sorted(students, key=lambda s: s['score'], reverse=True))  # 好
"""
[{'name': 'Lily', 'score': 98},
{'name': 'Stanley', 'score': 94},
{'name': 'Bob', 'score': 87},
{'name': 'Well', 'score': 85}]
"""

nums = [1,3,5,67,7,34,6,8,5,-12,-45,-234,232]
heapq.heapify(nums)
print(nums)
# [-234, -45, 1, 5, -12, 5, 6, 8, 67, 3, 7, 34, 232]
print(heapq.heappop(nums))  # heappop 返回的是序列中的第一个元素，也就是最小的一个元素
# -234


# 使用heapq编写优先级队列
import heapq


class PriorityQueue(object):
    def __init__(self):
        self._queue = []
        self._index = 0

    def push(self, item, priority):
        heapq.heappush(self._queue, (-priority, self._index, item))
        # 第一个参数是添加进的目标序列，
        # 第二个参数是将一个元组作为整体添加进序列，目的是为了方便比较，
        # 在priority相等的情况下，比较_index
        # priority为负数使得添加时按照优先级从大到小排序，因为堆排序的序列的第一个元素永远是最小的
        self._index += 1

    def pop(self):
        # 返回按照-priority 和 _index 排序后的第一个元素(是一个元组)的最后一个元素(item)
        return heapq.heappop(self._queue)[-1]

q = PriorityQueue()
q.push("bar", 2)
q.push("foo", 1)
q.push("gork", 3)
q.push("new", 1)

print(q.pop())
print(q.pop())
print(q.pop())
print(q.pop())
"""
gork  # 优先级最高
bar   # 优先级第二
foo   # 优先级与new相同，比较index，因为先加入，index比new小，所以排在前面
new
"""

参考资料：
　　Python Cookbook, 3rd edition, by David Beazley and Brian K. Jones (O’Reilly).