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)