commit

e/restaurant.py

@@ -0,0 +1,51 @@
+import math
+from queue import PriorityQueue
+
+
+def read_idx(x: str):
+    return int(x) - 1
+
+
+def inputs() -> list[str]:
+    return input().strip(" ").split(" ")
+
+
+V, E = map(int, inputs())
+s, t = map(read_idx, inputs())
+vw = list(map(int, inputs()))  # vertex weight
+graph = [[] for _ in range(V)]
+for _ in range(E):
+    u, v, w = map(read_idx, inputs())
+    w += 1
+    graph[u].append((v, w))
+    graph[v].append((u, w))
+
+newgraph = [[] for _ in range(V)]
+
+for u in range(V):
+    for v, w in graph[u]:
+        newgraph[u].append((v, vw[u] + w))
+
+
+def dijstra(g: list[list[tuple[int, int]]], s: int, t: int) -> float:
+    dist = [math.inf] * V
+    dist[s] = 0
+    q = PriorityQueue()
+    q.put((0, s))
+    while not q.empty():
+        d, u = q.get()
+        if d != dist[u]:
+            continue
+
+        # print(u, d)
+        for v, w in g[u]:
+            if dist[v] > dist[u] + w:
+                dist[v] = dist[u] + w
+                q.put((dist[v], v))
+
+    return dist[t]
+
+
+# print(vw, graph)
+# print(newgraph)
+print(dijstra(newgraph, s, t))