코드

import heapq
INF = int(1e9)

# v : 노드의 개수, e : 간선의 개수
v, e = map(int, input().split())

# graph는 무방향 그래프
graph = [[] for _ in range(v+1)]

# 그래프 정보 입력
for _ in range(e):
    a, b, c = map(int, input().split())
    
    # 무방향 그래프이므로 양쪽으로 다 넣어줘야 함
    graph[a].append((b, c))
    graph[b].append((a, c))
    
# 반드시 거쳐야 하는 두 개의 정점
nodes = list(map(int, input().split()))

# 시작은 1부터
start = 1

# 다익스트라 알고리즘
def dijstra(start):
    
    # 각각의 시작점에 따라 최단 거리 테이블이 달라지므로
    # 최단 거리 테이블을 메소드 안에서 선언하여 반환해줌
    dist_table = [INF for _ in range(v+1)]
    q = []
    heapq.heappush(q, (0, start))
    dist_table[start] = 0
    
    while q:
        dist, now = heapq.heappop(q)
        
        if dist < dist_table[now]:
            dist_table[now] = dist
            
        for i in graph[now]:
            cost = i[1] + dist
            
            if cost < dist_table[i[0]]:
                dist_table[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))
    
    # 최단 거리 테이블이 필요하므로 반환해줌
    return dist_table

# 다익스트라 수행 후 업데이트된 최단 거리 테이블은
# 파라미터로 받은 노드 번호에서 출발하여
# 각 테이블의 인덱스에 해당하는 번호를 가진 노드에 도착하기까지
# 필요한 최단 거리를 의미함!!!!
total = dijstra(start)       # 시작점(1)에서 다익스트라 수행
part1 = dijstra(nodes[0])    # 반드시 거쳐야 하는 노드 중 처음 노드에서 다익스트라 수행
part2 = dijstra(nodes[1])    # 반드시 거쳐야 하는 노드 중 두 번째 노드에서 다익스트라 수행

# start에서 시작하여 nodes[0]까지 가는 거리 
# +
# nodes[0]에서 시작하여 nodes[1]까지 가는 거리
# +
# nodes[1]에서 시작하여 v번(마지막)노드까지 가는 거리
pref_result = total[nodes[0]] + part1[nodes[1]] + part2[v]

# start에서 시작하여 nodes[1]까지 가는 거리 
# +
# nodes[1]에서 시작하여 nodes[0]까지 가는 거리
# +
# nodes[0]에서 시작하여 v번(마지막)노드까지 가는 거리
suf_result = total[nodes[1]] + part2[nodes[0] + part1[v]]

# 위의 두 경로 중 더 작은 값을 result로 선택한다!! => 이것이 곧 정답
result = min(pref_result, suf_result)

# result가 INF보다 작은 경우 result 출력
# 처음 result == INF로 조건을 분기했는데 99%에서 오답 발생
# 조건을 수정하니 정답으로 인정
if result < INF:
    print(result)
else:
    print(-1)

문제 해설

다익스트라 알고리즘의 응용 문제.

주어진 그래프에, 특정한 두 노드를 무조건 적으로 통과할 때, 마지막 번호의 노드로 도착하는 최단 거리를 구하라는 문제이다. 처음에 어떤 식으로 접근해야 할지 막막하였으나, 결국 두 노드를 의무적으로 통과하려면 아래와 같은 2가지 경우가 존재한다.

시작 노드를 1, 마지막 노드를 N, 두 노드를 v1, v2로 할 때

  • 1 -> v1 -> v2 -> N
  • 1 -> v2 -> v1 -> N

위 두 경우 중 더 작은 이동 비용이 발생하는 경우를 min()을 활용하여 출력시키면 정답이다.

다익스트라 알고리즘은 구현 자체가 생각보다 난이도가 있어 외울 수 있을 때까지 열심히 반복 숙달하여야겠다는 느낌이 강하게 든다. 반드시 마스터하여 내 것으로 만들어야겠다.

댓글남기기

참고