forked from ByteByteGoHq/coding-interview-patterns
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathshortest_path.py
More file actions
32 lines (30 loc) · 1.36 KB
/
shortest_path.py
File metadata and controls
32 lines (30 loc) · 1.36 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
from collections import defaultdict
import heapq
from typing import List
def shortest_path(n: int, edges: List[int], start: int) -> List[int]:
graph = defaultdict(list)
distances = [float('inf')] * n
distances[start] = 0
# Represent the graph as an adjacency list.
for u, v, w in edges:
graph[u].append((v, w))
graph[v].append((u, w))
min_heap = [(0, start)] # (distance, node)
# Use Dijkstra's algorithm to find the shortest path between the start node
# and all other nodes.
while min_heap:
curr_dist, curr_node = heapq.heappop(min_heap)
# If the current distance to this node is greater than the recorded
# distance, we've already found the shortest distance to this node.
if curr_dist > distances[curr_node]:
continue
# Update the distances of the neighboring nodes.
for neighbor, weight in graph[curr_node]:
neighbor_dist = curr_dist + weight
# Only update the distance if we find a shorter path to this
# neighbor.
if neighbor_dist < distances[neighbor]:
distances[neighbor] = neighbor_dist
heapq.heappush(min_heap, (neighbor_dist, neighbor))
# Convert all infinity values to -1, representing unreachable nodes.
return [-1 if dist == float('inf') else dist for dist in distances]