https://www.acmicpc.net/problem/1504
Solvedimport 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을 출력하도록 처리해줬다.