[다익스트라] heap을 사용한 다익스트라 예제 연습

코테에서 자주 나오는지는 모르겠지만 그래도 워낙에 유명한 알고리즘이다보니 공부하면 좋을 듯

일반적으로 최소 경로를 구하는 경우 모든 Edge의 가중치가 같으면 일반적인 BFS로 풀면 되지만
가중치가 다른 경우는 가장 최소 가중치 기준으로 방문을 수행하는 그리디 알고리즘이 필요하기 때문에 다익스트라가 요구된다.

여기서 다익스트라를 구현하는 많은 방법이 있지만 개인적으로 가장 파이썬에서 구현이 간단한 Heap을 사용한 방법을 익혀 코테에 써먹어보자.

예전에 자료구조 수업을 들을 때 얼핏 정리했던 코드를 가져와보았다. 
주로 이중 배열을 사용해서 구현을 하던데 나는 dictionary 형태로 배우다보니 그냥 지금도 이 폼을 사용한다.

import heapq

graph = {
    'A' : {'B':8, 'C':1, 'D':2, 'F':5},
    'B':{'A':8},
    'C': {'A':1,'B':5,'D':2},
    'D':{'A':2,'C':2,'E':3,'F':5},
    'E':{'D':3,'F':1},
    'F':{'A':5,'D':5,'E':1}
    }

def dijkstra(graph,start) :
    ### Start로부터 거리 값 초기 setting
    distance = { node : float('inf') for node in graph}
    distance[start] = 0
    queue = []
    heapq.heappush(queue,[distance[start],start])

    while queue :
        curr_dist, curr_dest = heapq.heappop(queue)

        ### 기존에 있는 거리보다 길다면, 볼 필요가 없음
        if distance[curr_dest] < curr_dist :
            continue

        for new_dest, new_dist in graph[curr_dest].items() :
            ## Edge Relaxation
            dist = curr_dist+new_dist
            if dist < distance[new_dest] :
                distance[new_dest] = dist
                heapq.heappush(queue,[dist,new_dest])
    return distance

print(dijkstra(graph,'A'))

일단 Case는 노드 당 하나의 경로만 공유하는 Undirected Graph 케이스 코드다. 물론 Directed는 그냥 그 값에 맞게 graph에 추가해주면 되는거라 크게 다르지는 않다.

사실 저 형태에서 크게 변하지 않는 거 같다. 저렇게 heap에 넣어서 중간에 distance를 갱신해주고 쭉쭉 이어가는형태로 BFS와 코드 구조가 비슷하지만 중간에 heap을 사용했다는 점이 일반적인 queue를 사용하는 BFS와는 차이가 있다.

그럼 이제 이걸 바탕으로 몇 문제를 풀이한 것을 정리해보자.
https://school.programmers.co.kr/learn/courses/30/lessons/12978

프로그래머스

 

위 문제의 경우 특별한 건 없고, 단 경로 별로 다른 가중치 값도 가질 수 있다는 점이 주의 포인트이다. 
하지만 이 경우 그냥 더 거리가 적은 값만 graph에 추가하면 됬으니 큰 문제는 없다.

import heapq as hp
from collections import defaultdict
def solution(N, road, K):
    graph = defaultdict(dict)
    distance = {node : float('inf') for node in range(1,N+1)} 
    ### Graph 설정할 때 양방향 setting으로 연결된 건 전부 연결해줌
    for element in road : 
        ### 동일한 경로가 2개 이상도 존재할 수 있음
        if element[1] in graph[element[0]]:
            ### 이 경우 최소 값만 포함
            if element[2] < graph[element[0]][element[1]] : 
                graph[element[0]][element[1]] = element[2]
        if element[0] in graph[element[1]]:
            if element[2] < graph[element[1]][element[0]] : 
                graph[element[1]][element[0]] = element[2]
        else : 
            graph[element[0]][element[1]] = element[2]
            graph[element[1]][element[0]] = element[2]

    queue = []
    distance[1] = 0
    hp.heappush(queue,[distance[1],1])
        
    while queue : 
        curr_dist, curr_dest = hp.heappop(queue)        
            
        if distance[curr_dest] < curr_dist : 
            continue
                
        for new_dest, new_dist in graph[curr_dest].items() : 
            dist = curr_dist+new_dist
            if dist < distance[new_dest] : 
                distance[new_dest] = dist
                hp.heappush(queue, [dist,new_dest])
    
    result = 0
    for node in distance :
        if distance[node] <= K : 
            result += 1
    print(distance)        
    return result

코드가 좀 더러운데, 나는 기본적으로 특별한 케이스가 아니면 그냥 단순한 undirected graph면 그냥 둘 다 이어주는 항을 만든다. 예를 들어 A-B가 연결되어 있으면 {A:{B:1,..}, ... B :{A:1,...}, ...} 이런식으로 2 항을 다 작성해준다.

다음 문제는 조금 더 특이한 문제인데, 여기서 주의해야 할 점이 바로,
"오고 / 가고" 가 포함된 경우다.

일반적으로 다익스트라는 시작점을 세팅하고 모든 노드에 대한 최단 경로를 구하는 식으로 동작하는데 (내가 배운 선에서는) 이를 활용하여 여러 번의 다익스트라를 구현하는 문제다. 
https://www.acmicpc.net/problem/1238

1238번: 파티

 위 문제의 경우 
먼저 다 같이 모이기로 한 지점을 시작 지점으로 하고
이 때 최단 경로 값들을 노드 별로 먼저 저장한 다음에
다음으로 모이기로 한 지점을 제외한 나머지 노드들에서 각각 다익스트라를 계산해
먼저 저장한 최단 경로 값들에 누적해서, 그 누적 값 중 가장 소요 시간이 오래된 경우를 return한다

즉 이 경우가 오고 가고 걸리는 시간이 가장 오래 걸리는 경우이기 때문이다.

from collections import defaultdict
import heapq as hp
N, M, dest = [int(x) for x in input().split()]

graph = defaultdict(dict)

for _ in range(M) : 
    start, end, cost = [int(x) for x in input().split()]
    ### 단 방향이므로 양 방향 처리를 해주지 않음
    graph[start][end] = cost

#print(graph)

def dijkstra(now) :
    d = {node: float('inf') for node in range(1,N+1)}
    d[now] = 0
    queue = []
    hp.heappush(queue,[d[now],now])

    while queue : 
        curr_dist, curr_dest = hp.heappop(queue)
        
        if d[curr_dest] < curr_dist : 
            continue
            
        for new_dest, new_dist in graph[curr_dest].items() : 
            dist = curr_dist+new_dist
            if dist < d[new_dest] : 
                d[new_dest] = dist
                hp.heappush(queue,[dist,new_dest])
    return d

ans = dijkstra(dest)
## {1: 1, 2: 0, 3: 3, 4: 7} 2에서 출발할 때
#print(ans)
for student in range(1,N+1):
    if student != dest :
        next = dijkstra(student)
        ans[student] += next[dest]

## {1: 5, 2: 0, 3: 9, 4: 10} 각 장소에서 2로 도착할 때의 경우 누적
#print(ans)
print(max(ans.values()))

#max_val = max(distance.values())
#print(max_val)

마지막으로 개인적으로 가장 어려웠는데 유형은 위 파티 문제와 유사
이 경우 두 개의 vertex를 지나는 최단 경로를 구하는 문제인데
처음에 최단 경로의 경우의 수를 구하는 문제인가 싶어 삽질하다가
앞의 파티 문제처럼 중간 노드를 활용한 다중 다익스트라를 통해 푸는 문제다.

즉 v1, v2를 반드시 거치고 1에서 N까지 가는 모든 경로를 탐색하려면

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

이 2가지 케이스를 분석해서 각각이 누적되는 다익스트라 값 중 최소 값을 리턴하면 된다

그리고 이 경우 백준에서 input()을 쓰면 시간 초과가 되므로 sys.stdin.readline()을 써야지만 통과됨

from collections import defaultdict
import heapq as hp
import sys

N, E = [int(x) for x in sys.stdin.readline().split()]

graph = defaultdict(dict)
for _ in range(E) : 
    start, end, cost = [int(x) for x in sys.stdin.readline().split()]
    graph[start][end] = cost
    graph[end][start] = cost

v1, v2 = [int(x) for x in input().split()]

def dijkstra(start) : 
    distance = {node : float('inf') for node in range(1,N+1)}
    distance[start] = 0
    queue = []
    hp.heappush(queue,[distance[start],start])

    while queue : 
        curr_dist, curr_dest = hp.heappop(queue)
    
        if distance[curr_dest] < curr_dist : 
            continue
        
        for new_dest, new_dist in graph[curr_dest].items() : 
            dist = curr_dist+new_dist 
            if dist < distance[new_dest] : 
                distance[new_dest] = dist
                hp.heappush(queue,[dist,new_dest])
    return distance

answer1 = dijkstra(v1)[1]
answer1 += dijkstra(v2)[v1]
answer1 += dijkstra(N)[v2]

answer2 = dijkstra(v2)[1]
answer2 += dijkstra(v1)[v2]
answer2 += dijkstra(N)[v1]

if answer1 == float('inf') :
    print(-1)
    exit(0)

print(min(answer1,answer2))