KK's blog

每天积累多一些

0%

Heap

Posted on Edited on Disqus:

算法思路：

最小堆可以维持堆顶元素为最小值。

应用：

  1. 求数组第k个大的数

Python代码：

1
2
3
4
5
6
7
8
9
def min_heap(self, nums: List[int], k: int) -> List[int]:
	if not nums:
		return []
	res = []
	for i in range(len(nums)):
		heappush(res, nums[i])
		if len(res) > k:
			heappop(res)
	return res
此法不用与堆顶元素比较，写法更简单，复杂度基本一样

max heap的话，入堆的数转负数，出堆后转为整数

1
2
3
4
5
6
7
8
def max_heap(self, nums: List[int], k: int) -> List[int]:
	res = []
	for i in range(len(nums)):
		heappush(res, -nums[i])
		if len(res) > k:
			heappop(res)
	res = [-n for n in res]
	return res

旧法不推荐

1
2
3
4
5
6
7
8
9
10
from heapq import heapreplace, heappush
def min_heap(self, nums: List[int], k: int) -> List[int]:
	res = []
	for i in range(len(nums)):
		if i < k:
			heappush(res, nums[i])
		elif nums[i] > res[0]:
			heapreplace(res, nums[i])

	return res

max heap的话，入堆的数转负数，跟堆顶比较的大于号不变，出堆后转为整数
注意第6行res[0]并没有负号，因为res已经是负数 

1
2
3
4
5
6
7
8
9
def max_heap(self, nums: List[int], k: int) -> List[int]:
	res = []
	for i in range(len(nums)):
		if i < k:
			heappush(res, -nums[i])
		elif -nums[i] > res[0]:
			heapreplace(res, -nums[i])
	res = [-n for n in res]
	return res

BFS + Heap

本质是图，求点或点的和或线段路径的最值 这是单源最短路径Dijkstra的典型应用。Dijkstra是每条边的权重(距离), 而不是节点个数最短(BFS模板)。
区别是

  1. 用BFS distance模板。queue变成heap每个点总weight的最小堆
  2. 遍历时候每个元素多一个weight
  3. visited不再是记录这个节点访问过没有，因为节点可以多次访问，目标是找到这个节点权重最小的路径。visited记录每个节点的最小权重(路径). visited的处理在neighbor的for循环中，比较权重是否最小，如果不是就跳过。 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    def bfs(self, graph, start, target) -> List[int]:
    	heap = [(0, k)] # weight, node
    	visited = {k: 0}
    	while heap:
    		weight, node  = heapq.heappop(heap)
    		for neighbor, _weight in graph[node]:
    			if neighbor in visited and visited[neighbor] <= weight + _weight:
    				continue
    			heapq.heappush(heap, (weight + _weight, neighbor))
    			visited[neighbor] = weight + _weight
    	return max(visited.values()) if len(visited) == n else -1
    BFS的distance模板(用于对比加强对BFS+Heap模板的理解) 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    def bfs_layer_v2(self, graph, start, target) -> int:
    	queue = deque([(start, 1)])
    	visited = {start}
    	while queue:
    		node, distance = queue.popleft()
    		if node == target:
    			return distance
    		for neighbor in graph[node]:
    			if neighbor in visited:
    				continue
    			queue.append((neighbor, distance + 1))
    			visited.add(neighbor)
    	return -1

算法分析：

时间复杂度为O(nlogk)，空间复杂度O(1)。

例子：

LeetCode 743 Network Delay Time 求从某一个点出发的所有能到达的点中的最短时间 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def networkDelayTime(self, times: List[List[int]], n: int, k: int) -> int:
	# build graph
	graph = collections.defaultdict(list)
	for e in times:
		graph[e[0]].append((e[1], e[2]))
	heap = [(0, k)] # weight, node
	visited = {k: 0}
	while heap:
		weight, node  = heapq.heappop(heap)
		for neighbor, _weight in graph[node]:
			if neighbor in visited and visited[neighbor] <= weight + _weight:
				continue
			heapq.heappush(heap, (weight + _weight, neighbor))
			visited[neighbor] = weight + _weight
	return max(visited.values()) if len(visited) == n else -1

LeetCode 787 Cheapest Flights Within K Stops
求只允许停k个站情况下，最便宜机票价格
weight在这里是每条边的价格 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
	graph = collections.defaultdict(list)
	for pair in flights:
		graph[pair[0]].append((pair[1], pair[2]))
	heap = ([(0, src, 0)]) # price, node_id, distance
	visited = {(src, 0): 0} # (node, distance): price
	while heap:
		p, node, dist = heapq.heappop(heap)
		if node == dst:
			return p
		if dist == k + 1:
			continue
		for neighbor, _price in graph[node]:
			if (neighbor, dist + 1) in visited and visited[(neighbor, dist + 1)] <= p + _price:
				continue
			heapq.heappush(heap, (p + _price, neighbor, dist + 1))
			visited[(neighbor, dist + 1)] = p + _price
	return -1

更简洁的写法 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
	graph = collections.defaultdict(list)
	for pair in flights:
		graph[pair[0]].append((pair[1], pair[2]))
	heap = ([(0, src, 0)]) # price, node_id, distance
	visited = {}
	while heap:
		p, node, dist = heapq.heappop(heap)
		if node == dst and dist <= k + 1:
			return p
		if node in visited and dist >= visited[node]:
			continue
		visited[node] = dist
		for neighbor, _price in graph[node]:
			heapq.heappush(heap, (p + _price, neighbor, dist + 1))
	return -1

Free mock interview