PythonでJOI難易度5を埋める #9

F - 船旅

考えたこと
最短経路を求める問題なので、クエリに従って辺を追加して、dijkstraする。無向グラフなので、両方向に辺を追加する。

from heapq import heappush, heappop

n, k = map(int,input().split())
info = [list(map(int,input().split())) for _ in range(k)]


INF = float('inf')
def dijkstra(s, n, m):
    dist = [INF] * n
    hq = [(0, s)]
    dist[s] = 0
    seen = [False] * n
    while hq:
        v = heappop(hq)[1]
        seen[v] = True
        for to, cost in m[v]:
            if not seen[to] and dist[v] + cost < dist[to]:
                dist[to] = dist[v] + cost
                heappush(hq, (dist[to], to))
    return dist


maps = [[] for _ in range(n)]

for i in range(k):
    if info[i][0] == 0:
        res = dijkstra(info[i][1]-1, n, maps)
        ans = res[info[i][2]-1]
        if ans == float('inf'):
            print(-1)
        else:
            print(ans)
    if info[i][0] == 1:
        maps[info[i][1]-1].append((info[i][2]-1, info[i][3]))
        maps[info[i][2]-1].append((info[i][1]-1, info[i][3]))