https://www.acmicpc.net/problem/1504
Solved
import sys, heapq
input = sys.stdin.readline
N, E = map(int, input().rstrip().split())
arr = [[] for _ in range(N+1)]
for _ in range(E):
u, v, w = map(int, input().rstrip().split())
arr[u].append((v, w)) # u -> v : w
arr[v].append((u, w)) # v -> u : w (무방향 가중치 그래프)
V1, V2 = map(int, input().rstrip().split())
def dijkstra(arr, start):
costs = [float('inf')] * (N+1)
pq = []
heapq.heappush(pq, (0, start))
costs[start] = 0
while pq:
cur_cost, cur_node = heapq.heappop(pq)
if costs[cur_node] < cur_cost:
continue
for next_node, cost in arr[cur_node]:
next_cost = cost + cur_cost
if costs[next_node] > next_cost:
costs[next_node] = next_cost
heapq.heappush(pq, (next_cost, next_node))
return costs
costs_from_1 = dijkstra(arr, 1) # 1에서 다른 모든 정점까지의 최단거리
costs_from_V1 = dijkstra(arr, V1) # V1에서 다른 모든 정점까지의 최단거리
costs_from_V2 = dijkstra(arr, V2) # V2에서 다른 모든 정점까지의 최단거리
# 경로 1 : 1 -> V1 -> V2 -> N
path1 = costs_from_1[V1] + costs_from_V1[V2] + costs_from_V2[N]
# 경로 2 : 1 -> V2 -> V1 -> N
path2 = costs_from_1[V2] + costs_from_V2[V1] + costs_from_V1[N]
min_cost = min(path1, path2)
print(min_cost if min_cost != float('inf') else -1)
다익스트라 로직의 경우 기존 다익스트라 코드와 동일하게 작성해줬다.
다른 다익스트라 문제와 크게 다른 점이 있다면, 다익스트라 함수를 여러번 호출해 경로끼리 비교를 해야 한다는 점이다.
문제 입력에 마지막으로 주어지는 V1
과 V2
를 반드시 지나서 정점 N
에 가야하므로, 이 부분에 대한 로직을 추가적으로 생각해야 한다.
처음에는 다익스트라 함수 내부에서 V1
과 V2
를 지나는 경로를 끼워넣을까 생각했었다.
다시 생각해보면, 다익스트라 함수에서는 arr
과 start
를 인자로 받고 있다. 그 이유는 start
정점을 시작으로 모든 정점에 도착하기까지의 최단 경로를 계산해 리턴하기 때문이다.
그럼! 이 문제에서는 정점 1에서 시작해 정점 N에 도달하는 과정에서 V1과 V2를 반드시 거처야 하므로 나올 수 있는 경로는
1. 1 ➡️ V1 ➡️ V2 ➡️ N
2. 1 ➡️ V2 ➡️ V1 ➡️ N
위 두가지이다.
costs_from_1 = dijkstra(arr, 1) # 1에서 다른 모든 정점까지의 최단거리
costs_from_V1 = dijkstra(arr, V1) # V1에서 다른 모든 정점까지의 최단거리
costs_from_V2 = dijkstra(arr, V2) # V2에서 다른 모든 정점까지의 최단거리
위와 같이 1에서, V1에서, V2에서 모든 정점으로 가는 데 걸리는 최단거리를 계산해준다.
아! 이 때 주의해야 할 점은
for _ in range(E):
u, v, w = map(int, input().rstrip().split())
arr[u].append((v, w)) # u -> v : w
arr[v].append((u, w)) # v -> u : w (무방향 가중치 그래프)
문제에서 "방향성이 없는 가중치 그래프"라고 명시해두었으니,
u ➡️ v : w
v ➡️ u : w
위와 같이 두방향으로 향하는 경로를 모두 고려해줘야 한다.
1, V1, V2에서 모든 경로로 뻗어가는 비용은 다 계산했으니,
1. 1 ➡️ V1 ➡️ V2 ➡️ N
2. 1 ➡️ V2 ➡️ V1 ➡️ N
이 두가지 경로로 가는 비용 중 더 적은 비용이 나온 경로를 채택하는 방법을 생각해보자.
# 경로 1 : 1 -> V1 -> V2 -> N
path1 = costs_from_1[V1] + costs_from_V1[V2] + costs_from_V2[N]
# 경로 2 : 1 -> V2 -> V1 -> N
path2 = costs_from_1[V2] + costs_from_V2[V1] + costs_from_V1[N]
경로 1에 대해 설명하자면,
1에서 시작해 V1에 도달한 비용
+ (1 -> V1으로 왔으므로 현재 위치는 V1) V1에서 시작해 V2에 도달한 비용
+ (V1 -> V2로 왔으므로 현재 위치는 V2) + V2에서 시작해 N에 도달한 비용
(현재 위치는 V2 -> N(도착점)으로 이동)
위와 같은 흐름이라고 할 수 있다.
경로 2 역시 동일한 방법으로 계산해준 후 최솟값을 갱신해주면 된다.
min_cost = min(path1, path2)
print(min_cost if min_cost != float('inf') else -1)
경로가 없을 경우 -1을 출력해줘야 하니까!
costs의 초기화 값이었던 float('inf')
일 경우 -1을 출력하도록 처리해줬다.