heapq模块提供了堆队列算法的实现，也称为优先队列算法

堆的创建

我们可以可以使用两种方式来创建堆。一个是使用 heapq.heappush 的方式来向堆中不断的填充元素。另外一种方式是使用 heapq.heapify() 来将list转化为heap结构。

For example

1
2
3
4
5
6
7
8
9
import heapq
l = [1,2,3,4]
heap = []
# solution 1
for val in l:
	heapq.heappush(heap,val) 
# solution 2
heapq.heapify(l)

堆的访问

弹出最小值

可以使用函数 heapq.heappop(heap) 来弹出堆中的最小值。传入的参数类型为堆

For example

1
2
3
4
import heapq
nums = [1,3,2,5,4]
heapq.heapify(nums)
print(heapq.heappop(nums)) # out: 1

弹出当前最小值并添加一个新的元素

可以使用函数 heapq.heapreplace 来在弹出最小值后再添加一个元素。

For example

1
2
3
4
5
6
import heapq
nums = [1,3,2,5,4]
heapq.heapify(nums)
heapq.heapreplace(nums,1.5) # will pop 1 and add 1.5
# 查看当前堆的最小值
print(nums[0]) # out: 1.5

使用 nums[0] 可以快捷的访问当前堆的最小值

添加一个新元素并弹出当前堆的最小值

可以使用函数 heapq.heappushpop(heap,item)来实现将一个元素放入堆中后，返回当前堆的最小值。

For Example

1
2
3
4
5
6
7
8
import heapq
nums = [1,3,2,5,4]
heapq.heapify(nums)
print(heapq.heappushpop(6,nums))
print(nums)
# out:
# 1
# [2, 3, 6, 5, 4]

获取当前堆的最小或者最大值

如果需要获取当前堆前n个的最小值或者最大值，可以使用 heapq.nlargest() or heapq.nsmalles()

For example

1
2
3
4
5
6
7
8
import heapq
nums = [1,3,2,5,4]
heapq.heapify(nums)
print(heapq.nlargest(2,nums))
print(heapq.nsmallest(2,nums))
# out:
#[5, 4]
#[1, 2]

应用

Leetcode 703题 数据流中的第 K 大元素

题目要求：

设计一个找到数据流中第 k 大元素的类（class）。注意是排序后的第 k 大元素，不是第 k 个不同的元素。

请实现 KthLargest 类：

KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。
int add(int val) 将 val 插入数据流 nums 后，返回当前数据流中第 k 大的元素。

示例：

1
2
3
4
5
6
7
8
9
10
11
12
13
输入：
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
输出：
[null, 4, 5, 5, 8, 8]

解释：
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3);   // return 4
kthLargest.add(5);   // return 5
kthLargest.add(10);  // return 5
kthLargest.add(9);   // return 8
kthLargest.add(4);   // return 8

解答

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import heapq
class KthLargest:
	def __init__(self,k:int,nums:List[int]):
		self.pool = nums
		heapq.heapify(self.pool)
		self.k = k
		while len(self.pool) > k:
			heapq.heappop(self.pool)
	def add(self,val:int) -> int:
		if len(self.pool) < self.k:
			heapq.heappush(self.pool,val)
		elif val > self.pool[0]:
			heapq.heapreplace(self.pool,val)
		return self.pool[0]

解答思路

维持一个最大为K的最小堆。每次添加元素，如果元素大于堆的最小值，则将当前堆的最小值pop,然后将元素push到堆中。最后返回元素添加后的最小值。如果当前堆的大小小于K,则直接将此元素添加到堆中。然后返回堆的最小值。

参考

[1] https://leetcode-cn.com/problems/kth-largest-element-in-a-stream

[2] https://docs.python.org/zh-cn/3/library/heapq.htm

[3] https://www.jianshu.com/p/801318c77ab5