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

Dijkstra 문제

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

1504번: 특정한 최단 경로

 

💡 풀이코드 1  (성공)

import sys
import heapq
import copy
node, edge = map(int, sys.stdin.readline().split())
graph = [[] for _ in range(node + 1)]
visited = [[False] for _ in range(node + 1)]

for i in range(edge):
    u, v, w = map(int, sys.stdin.readline().split())
    graph[u].append((v, w))
    graph[v].append((u, w))
first, second = map(int, sys.stdin.readline().split())

def dijkstra(graph, start, end):
    heap = [(0, start)]
    distance = [sys.maxsize] * (node + 1)
    distance[start] = 0

    while heap:
        curr_dist, u = heapq.heappop(heap)
        if distance[u] < curr_dist: continue
        for v, weight in graph[u]:
            if distance[v] > distance[u] + weight:
                distance[v] = distance[u] + weight
                heapq.heappush(heap, (distance[v], v))
    return distance[end]

answer = sys.maxsize
l = [1, first, first, second, second, node]
for i in range(4):
    cnt = 0
    for j in range(3):
        j *= 2
        cnt += dijkstra(graph, l[j], l[j + 1])
        if i == 2 and j ==0:
            cnt *= 2
        elif i == 3 and j == 0:
            cnt *= 2
    answer = min(answer, cnt)
    if i == 0:
        l = [1, second, second, first, first, node]
    elif i == 1:
        l = [1, second, 1, first, first, node]
    elif i == 2:   
        l = [1, first, 1, second, second, node]     

print(-1 if answer == sys.maxsize else answer)
# dijkstra 4번
# 1->2, 2->3, 3->4
# 1->3, 3->2, 2->4

# 1->2 1->3 2->4
# 1->2 1->3 3->4

 

정답을 맞추긴 했지만, 위의 코드는 너무 하드코딩스럽게 짰다.

우리는 최단경로를 구할 때, 마지막에 입력으로 받은 두 정점을 무조건 지나게끔 !! 최단경로를 구해야 한다. 

따라서, 1번 정점에서 N번 정점까지의 최단경로를 다이렉트로 바로 구하는 것이 아니라, 다음과 같이 최단경로를 나눠서 구한 후 dijkstra 함수 돌린 값을 모두 더해주면 된다.

 

예를 들어, 2번 정점과 3번 정점을 무조건 지나야 한다고 가정해보자. 

(원래) : 1 -> N까지 dijkstra 함수 1번 호출

(변경 후) : 1 -> 2, 2 -> 3, 3 -> 4 까지 dijkstra 함수 3번 호출

                 1 -> 3, 3 -> 2, 2 -> 4

 -> 중간에 2번 /3번 정점 중 어떤 정점을 먼저 거치냐로 경로가 나뉠 수도 있기 때문에 이렇게 두가지로 나눠서 구한 후, 최단경로를 선택해주면 된다.

즉, 이렇게만 하면 총 dijkstra 함수를 6번 호출하는 것이다.

 

하지만, 나는 문제 설명에서 이 문제가 양방향 그래프 기준이라는 글씨를 못보고, 단방향 그래프 기준으로 입력을 받아 풀었었다.

따라서, 나는 이 문제를 단방향 그래프 기준으로 놓고 풀었기 때문에 위의 2가지 경우에서 추가로 2가지 경우를 더 생각해주어야 했다.

(추가 후) : 1 -> 2, 2 -> 3, 3 -> 4 까지 dijkstra 함수 3번 호출

                 1 -> 3, 3 -> 2, 2 -> 4

                 1 -> 3, 1 -> 2, 2 -> 4 

                 1 -> 2, 1 -> 3, 3 -> 4

위에서, 3번째와 같은 경우에는 1 -> 3, 3 -> 1로 왕복해서 다시 온 후, 1에서 2를 거치는 것이기 때문에 1 -> 3의 가중치 값을 2배 해준다.

4번째 경우도 마찬가지로 1 -> 2, 2 -> 1로 왕복해서 다시 온 후, 1에서 3을 거치는 것이기 때문에 1 -> 2의 가중치 값을 2배 해준다.

 

하지만, 내가 위의 2가지 경우를 추가해줬기 때문에 문제가 풀린 것이 아니라, 이 문제가 양방향 그래프라는 사실을 깨달은 후 graph 입력 받는 부분을 고쳐줬기 때문에 이 문제가 풀렸던 것이다..

 

💡 풀이코드 2 (성공,  코드 리팩토링)

import sys
import heapq
node, edge = map(int, sys.stdin.readline().split())
graph = [[] for _ in range(node + 1)]

for i in range(edge):
    u, v, w = map(int, sys.stdin.readline().split())
    graph[u].append((v, w))
    graph[v].append((u, w))
first, second = map(int, sys.stdin.readline().split())

def dijkstra(graph, start, end):
    heap = [(0, start)]
    distance = [sys.maxsize] * (node + 1)
    distance[start] = 0

    while heap:
        curr_dist, u = heapq.heappop(heap)
        if distance[u] < curr_dist: continue
        for v, weight in graph[u]:
            if distance[v] > distance[u] + weight:
                distance[v] = distance[u] + weight
                heapq.heappush(heap, (distance[v], v))
    return distance[end]

answer = sys.maxsize
l = [1, first, second, node]
for i in range(2):
    cnt = 0
    for j in range(3):
        cnt += dijkstra(graph, l[j], l[j + 1])
    answer = min(answer, cnt)
    l = [1, second, first, node]
   

print(-1 if answer == sys.maxsize else answer)
# dijkstra 2번
# 1->2, 2->3, 3->4
# 1->3, 3->2, 2->4

 

위에 추가로 넣었던 2가지 경우는 필요가 없다.

 

(재변경 후) : 1 -> 2, 2 -> 3, 3 -> 4 까지 dijkstra 함수 3번 호출

                 1 -> 3, 3 -> 2, 2 -> 4

                 1 -> 3, 3 -> 1, 1 -> 2, 2 -> 4 

                 1 -> 2, 2 -> 1, 1 -> 3, 3 -> 4

 

2 -> 1, 1 -> 3이 결국엔 2 -> 3이다. 즉, 4번째 경우가 1번째 경우와 같다.

3 -> 1, 1 -> 2이 결국엔 3 -> 2이다. 즉, 3번째 경우가 2번째 경우와 같다.

 

따라서, 단방향 그래프인줄 알고, 추가로 넣어줬던 2가지 경우를 다시 빼주었다. 1번째 경우와 같이 loop를 돌며 dijkstra를 3번 호출한 이후엔 first와 second 변수의 자리를 바꿔 2번째 경우와 같이 함수 호출을 할 수 있도록 해주었다.

 

 

결론은 코드 리팩토링 후, 메모리와 시간이 모두 줄었다 !

 

*참고

https://www.acmicpc.net/board/view/124591

글 읽기 - 3%, 10%, 75% 반례입니다