[백준] 1504 특정한 최단 경로 (python 파이썬)

https://www.acmicpc.net/problem/1504

1504번: 특정한 최단 경로

---

아이디어

1. 가중치가 있는 그래프

• 다익스트라 아니면 플로이드 두 개를 생각했다.
• 두정점 v1, v2를 지나는 최단경로는 각각의 최단경로들을 더한 것과 같으므로 다익스트라를 여러 번 사용하면 된다.
• '1 -> N' = '1 -> v1' + 'v1 -> v2' + 'v2 ->N' 같이 계산해서 두 정점을 지나는 최단경로를 구한다.
• 루트는 1 -> v1 -> v2 -> N, 1 -> v2 -> v1 -> N  두 가지를 계산해 비교하면 작은 값을 출력하면 된다.

 

---

전체 코드

import sys, heapq
INF = int(1e9)

N, E = map(int,input().split())

graph = [[] for _ in range(N + 1)]

for _ in range(E):
    a, b, c = map(int,sys.stdin.readline().rstrip().split())

    graph[a].append([b, c])
    graph[b].append([a, c])

v1, v2 = map(int,input().split())

def dijkstra(start):
    que = []
    heapq.heappush(que, (0, start))
    D[start] = 0

    while que:
        dist, now = heapq.heappop(que)
        if D[now] < dist:
            continue

        for i in graph[now]:
            cost = dist + i[1]
            if cost < D[i[0]]:
                D[i[0]] = cost
                heapq.heappush(que, (cost,i[0]))

root_A, root_B = 0, 0

D = [INF] * (N + 1)
dijkstra(1) # 1 -> v1, 1 -> v2
root_A += D[v1]
root_B += D[v2]

D = [INF] * (N + 1)
dijkstra(v1) # v1 -> v2, v1 -> N
root_A += D[v2]
root_B += D[N]

D = [INF] * (N + 1)
dijkstra(v2) # v2 -> N, v2 -> v1
root_A += D[N]
root_B += D[v1]

ans = min(root_A, root_B)

if ans >= INF:
    print(-1)
else:
    print(ans)

---

코드 설명

import sys, heapq
INF = int(1e9)

N, E = map(int,input().split())

graph = [[] for _ in range(N + 1)]

for _ in range(E):
    a, b, c = map(int,sys.stdin.readline().rstrip().split())

    graph[a].append([b, c])
    graph[b].append([a, c])

v1, v2 = map(int,input().split())

그래프의 연결 관계는 인접 리스트 형태로 입력받아준다.

 

def dijkstra(start):
    que = []
    heapq.heappush(que, (0, start))
    D[start] = 0

    while que:
        dist, now = heapq.heappop(que)
        if D[now] < dist:
            continue

        for i in graph[now]:
            cost = dist + i[1]
            if cost < D[i[0]]:
                D[i[0]] = cost
                heapq.heappush(que, (cost,i[0]))

데이크스트라 코드는 힙을 사용해서 구현했다. 

데이크스트라에 대한 내용은 넘어가겠다.

필요한 분들은 구글 검색을 사용해주세요.

 

root_A, root_B = 0, 0

D = [INF] * (N + 1)
dijkstra(1) # 1 -> v1, 1 -> v2
root_A += D[v1]
root_B += D[v2]

D = [INF] * (N + 1)
dijkstra(v1) # v1 -> v2, v1 -> N
root_A += D[v2]
root_B += D[N]

D = [INF] * (N + 1)
dijkstra(v2) # v2 -> N, v2 -> v1
root_A += D[N]
root_B += D[v1]

ans = min(root_A, root_B)

if ans >= INF:
    print(-1)
else:
    print(ans)

문제에서 두 정점을 지나가는 방법은 

1 -> v1 -> v2 -> N, 1 -> v2 -> v1 -> N(root_A, root_B) 이렇게 두 가지 방법이 존재한다.

한 번에 최단거리를 구하는 게 아닌 정점과 정점 사이 거리를 각각 구해 더해주는 방식으로 최단거리를 구하면 된다.

 

문제에서 이미 이동했던 정점이나 간선으로 다시 이동할 수 있다고 나와있으므로

데이크스트라를 여러 번 써서 최단거리를 구해도 괜찮다.

 

그래서 정점 1, v1, v2 총 3번 데이크스트라 함수를 실행해서 root_A와 root_B 값에 더해준다.

이후 root_A와 root_B 두 값 중 작은 값을 ans 변수에 저장한다.

 

문제에서 경로가 존재하지 않는 경우는 -1을 출력하라고 나와있다.

경로가 존재하지 않다는 얘기는 ans에 INF값이 더해졌다는 얘기다.

왜냐 INF값이 더해졌다는 건 어떤 노드와 노드 사이가 이어져있지 않다는 걸 뜻한다.

그래서 ans가 INF보다 크면 -1을 아니면 그냥 ans값을 출력해주면 된다.