求有向圖兩節點間的最短路，需要支持增加邊和多次詢問，Dijkstra
1 #2642
2 #Runtime: 566 ms
3 #Memory: 17.3 MB
4
5 class Graph(object):
6
7 def __init__(self, n, edges):
8 """
9 :type n: int
10 :type edges: List[List[int]]
11 """
12 self.adj = [[] for _ in xrange(n)]
13 for e in edges:
14 self.adj[e[0]].append((e[1], e[2]))
15
16
17 def addEdge(self, edge):
18 """
19 :type edge: List[int]
20 :rtype: None
21 """
22 self.adj[edge[0]].append((edge[1], edge[2]))
23
24
25 def shortestPath(self, node1, node2):
26 """
27 :type node1: int
28 :type node2: int
29 :rtype: int
30 """
31 return self.dijkstra(node1, node2)
32
33
34 def dijkstra(self, st, ed):
35 n = len(self.adj)
36 dis = [float('inf')] * n
37 dis[st] = 0
38 q = [(0, st)]
39 while q:
40 cur_cost, cur_node = heapq.heappop(q)
41 if cur_cost > dis[cur_node]:
42 continue
43 if cur_node == ed:
44 return cur_cost
45 for e in self.adj[cur_node]:
46 node, l = e
47 tp_cost = l + dis[cur_node]
48 if dis[node] > tp_cost:
49 dis[node] = tp_cost
50 heapq.heappush(q, (tp_cost, node))
51 return -1 if dis[ed] == float('inf') else dis[ed]
52
53
54
55 # Your Graph object will be instantiated and called as such:
56 # obj = Graph(n, edges)
57 # obj.addEdge(edge)
58 # param_2 = obj.shortestPath(node1,node2)