Python优先级队列
最近刷题做到K个有序链表的合并,用到优先级队列,对优先级队列做个记录
优先级队列
优先级队列是一种容器型数据结构,它能管理一队记录,并按照排序字段(例如一个数字类型的权重值)为其排序。由于是排序的,所以在优先级队列中你可以快速获取到最大的和最小的值。
优先级队列使用堆来实现,在Python中有heapq 和 PriorityQueue 两种实现.
heapq
heapq 默认采用最小堆实现
import heapq
arr= [1,3,4,2]
print("原数组:{0}".format(arr))
# 将给定的列表转化为最小堆,线性时间
heapq.heapify(arr)
print("最小堆数组:{0}".format(arr))
# 插入元素
heapq.heappush(arr, 1)
print("插入新元素后:{0}".format(arr))
# 弹出最小元素
item0 = heapq.heappop(arr)
print("弹出的元素后:{0}".format(arr))
# 返回最小元素
item1 = arr[0]
print("获取最小元素的值:{0}".format(item1))
# 弹出最小元素,并插入一个新的元素,相当于先 heappop, 再 heappush
item2 = heapq.heapreplace(arr, -2)
最大堆:
import heapq
arr = [1,4,2,3]
print("原数组:{0}".format(arr))
# 将给定的列表转化为最大堆,线性时间
heapq._heapify_max(arr)
print("最大堆数组:{0}".format(arr))
# 弹出最大元素
item0 = heapq._heappop_max(arr)
print("弹出的元素为:{0}".format(item0))
print("弹出的元素后:{0}".format(arr))
# 弹出最大元素,并插入一个新的元素
item1 = heapq._heapreplace_max(arr, 9)
print("弹出的元素为:{0}".format(item1))
print("现在的堆结构为:{0}".format(arr))
自定义比较函数
当使用元祖、字典等复杂结构时进行优先排列时,会先根据第一个元素进行比较、第一个元素相等再进行第二个元素的比较。如果自己实现类要使用 heapq, 需要复写 lt 函数,比如:
class ListNode:
def __init__(self, val=0, next=None):
self.val = val
self.next = next
def __lt__(self, other):
return self.val < other.val
返回最大、最小的K个元素
heapq 自带函数获取可迭代对象中最大、最小的K个元素
import heapq
arr = [1,2,3,45,0,6,6,4]
heapq.heapify(arr)
heapq.nlargest(3, arr)
heapq.nsmallest(2, arr)