백준 1504번: 특정한 최단 경로 (G4) - Python 풀이

문제 설명

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

요약) 1번 정점에서 N번 정점으로 최단 거리로 이동해야 하는데, 주어진 두 정점도 꼭 거쳐야 합니다.
즉, 1번 정점, 주어진 정점 v1, 주어진 정점 v2, N번 정점 → 총 4개의 정점을 거치는 최단 거리를 출력해야 합니다.

풀이 과정

그렇다면 가능한 경로는 1번 → v1 → v2 → N번 아니면 1번 → v2 → v1 → N번밖에 없을 것입니다.

1트 틀린 코드

import heapq

dest, lines_num = map(int, input().split())

INF = 1e8
graph = [[] for _ in range(dest + 1)]

for _ in range(lines_num):
    node1, node2, cost = map(int, input().split())
    graph[node1].append((node2, cost))
    graph[node2].append((node1, cost))

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

def dijkstra(start):
    distance = [INF for _ in range(dest + 1)]
    q = []
    distance[start] = 0
    heapq.heappush(q, (0, start))

    while q:
        shortest_distance, index = heapq.heappop(q)
    
        for next_node, dist_to_next in graph[index]:
            if distance[next_node] > shortest_distance + dist_to_next:
                distance[next_node] = shortest_distance + dist_to_next
                heapq.heappush(q, (shortest_distance + dist_to_next, next_node))
    return distance
    
shortest_distances_start_1 = dijkstra(1)
shortest_distances_start_v1 = dijkstra(v1)
shortest_distances_start_v2 = dijkstra(v2)

v1_v2_path = shortest_distances_start_1[v1] + shortest_distances_start_v1[v2] + shortest_distances_start_v2[dest]
v2_v1_path = shortest_distances_start_1[v2] + shortest_distances_start_v2[v1] + shortest_distances_start_v1[dest]

if shortest_distances_start_1[v1] == INF or shortest_distances_start_1[v2] == INF:
    print(-1)
else:
    print(min(v1_v2_path, v2_v1_path))

왜 틀렸을까요.

맨 하단에 보시면 v1과 v2 값이 INF라면 즉, 1에서부터 시작하는 다익스트라 알고리즘을 한 번 돌린 후에도 v1과 v2까지의 값이 무한이라면, 1번과 v1 또는 v2가 연결된 간선이 없다는 뜻이기 때문에 최단거리를 구할 수 없어서 -1을 출력했습니다.

그러나 반례 중에 1번과 N번이 간선으로 연결되어 있지 않은 경우가 있었기에 이 경우를 고려하려면 if 문에 shortest_distances_start_1[dest] == INF 조건을 추가해야 합니다.

정답 코드

import heapq

# dest : 정점의 총 개수이자 최종 목적지 정점 번호
# lines_num : 간선의 총 개수
dest, lines_num = map(int, input().split())

# 최댓값
INF = 1e8
# 정점과 간선의 정보를 저장할 graph
graph = [[] for _ in range(dest + 1)]

# 간선의 개수만큼 돌면서 그래프 값을 채운다
for _ in range(lines_num):
    node1, node2, cost = map(int, input().split())
    # 양방향 그래프이기 때문에 node1이 시작점인 경우와 node2가 시작점인 경우 모두에 추가한다
    graph[node1].append((node2, cost))
    graph[node2].append((node1, cost))

# 꼭 지나야 하는 정점 2개
v1, v2 = map(int, input().split())

# start에서 각 정점까지의 최단 거리를 구하는 다익스트라 알고리즘
def dijkstra(start):
    # distance : start에서 각 정점까지의 최단거리를 저장할 1차원 배열
    # 정점의 총 개수만큼 최댓값으로 초기화
    distance = [INF for _ in range(dest + 1)]
    # 큐 초기화
    q = []
    # 자기 자신까지의 거리는 0
    distance[start] = 0
    # 최소힙에 시작 정점을 push
    heapq.heappush(q, (0, start))
	
    # 큐에 값이 존재하지 않을 때까지 반복
    while q:
        # 최소힙에서 pop
        # shortest_distance : 정점 중 최단 거리
        # index : 그 정점의 번호
        shortest_distance, index = heapq.heappop(q)
		    
        # 최단 거리를 가진 정점과 연결된 정점들을 확인
        # next_node : 연결된 정점의 번호
        # dist_to_next : index에서 연결된 정점 사이의 거리
        for next_node, dist_to_next in graph[index]:
            # 현재 next까지의 최단 거리보다 index를 경유하여 next로 가는 거리 값이 더 작으면
            if distance[next_node] > shortest_distance + dist_to_next:
                # 그 값으로 next까지의 최단 거리를 업데이트
                distance[next_node] = shortest_distance + dist_to_next
                # 최소힙에 push
                heapq.heappush(q, (shortest_distance + dist_to_next, next_node))
                
    # 최종적으로 start에서 각 정점까지의 최단 거리 배열을 반환
    return distance
    
# 1에서 각 정점까지의 최단 거리 배열
shortest_distances_start_1 = dijkstra(1)
# v1에서 각 정점까지의 최단 거리 배열
shortest_distances_start_v1 = dijkstra(v1)
# v2에서 각 정점까지의 최단 거리 배열
shortest_distances_start_v2 = dijkstra(v2)

# 1 -> v1 -> v2 -> dest의 순서대로 이동한 최단 거리
# shortest_distances_start_1[v1] : 1 -> v1 까지의 최단 거리
# shortest_distances_start_v1[v2] : v1 -> v2 까지의 최단 거리
# shortest_distances_start_v2[dest] : v2 -> dest 까지의 최단 거리
v1_v2_path = shortest_distances_start_1[v1] + shortest_distances_start_v1[v2] + shortest_distances_start_v2[dest]
# 1 -> v2 -> v1 -> dest의 순서대로 이동한 최단 거리
v2_v1_path = shortest_distances_start_1[v2] + shortest_distances_start_v2[v1] + shortest_distances_start_v1[dest]

# 1에서 v1 또는 v2 또는 dest 정점까지의 최단 거리가 INF로 나온다면
if shortest_distances_start_1[v1] == INF or shortest_distances_start_1[v2] == INF or shortest_distances_start_1[dest] == INF:
    # 1과 이어지는 간선이 없다는 뜻이기 때문에 -1을 출력
    print(-1)
else:
    # 두 경로 중 더 최솟값이 정답
    print(min(v1_v2_path, v2_v1_path))

캬 제가 봐도 주석이 너무 상세합니다. LGTM.

이상입니다.