[Day+16] 백준 특정 거리의 도시 찾기 (다익스트라 알고리즘)

오늘은 하루종일 다익스트라 알고리즘에 대해서 풀었다.

다익스트라 알고리즘에 대해서는 저번에 포스팅을 했으니 바로 문제 풀이로 넘어가겟다.

 

전체 코드

"""  
문제
어떤 나라에는 1번부터 N번까지의 도시와 M개의 단방향 도로가 존재한다. 모든 도로의 거리는 1이다.

이 때 특정한 도시 X로부터 출발하여 도달할 수 있는 모든 도시 중에서, 최단 거리가 정확히 K인 모든 도시들의 번호를 출력하는 프로그램을 작성하시오. 또한 출발 도시 X에서 출발 도시 X로 가는 최단 거리는 항상 0이라고 가정한다.

예를 들어 N=4, K=2, X=1일 때 다음과 같이 그래프가 구성되어 있다고 가정하자.



이 때 1번 도시에서 출발하여 도달할 수 있는 도시 중에서, 최단 거리가 2인 도시는 4번 도시 뿐이다.  2번과 3번 도시의 경우, 최단 거리가 1이기 때문에 출력하지 않는다.

입력
첫째 줄에 도시의 개수 N, 도로의 개수 M, 거리 정보 K, 출발 도시의 번호 X가 주어진다. (2 ≤ N ≤ 300,000, 1 ≤ M ≤ 1,000,000, 1 ≤ K ≤ 300,000, 1 ≤ X ≤ N) 둘째 줄부터 M개의 줄에 걸쳐서 두 개의 자연수 A, B가 공백을 기준으로 구분되어 주어진다. 이는 A번 도시에서 B번 도시로 이동하는 단방향 도로가 존재한다는 의미다. (1 ≤ A, B ≤ N) 단, A와 B는 서로 다른 자연수이다.

출력
X로부터 출발하여 도달할 수 있는 도시 중에서, 최단 거리가 K인 모든 도시의 번호를 한 줄에 하나씩 오름차순으로 출력한다.

이 때 도달할 수 있는 도시 중에서, 최단 거리가 K인 도시가 하나도 존재하지 않으면 -1을 출력한다.
"""
import heapq

def dijkstra(graph, start):
    distances=[float('inf')]*len(graph)
    distances[start-1]=0

    queue = [(0,start)]

    while queue :
        curr_distance, curr_node = heapq.heappop(queue)
        if curr_distance > distances[curr_node] :
            continue

        for neighbor in range(len(graph)):
            weight = graph[start][neighbor]
            
            if weight != 0 :
                distance = curr_distance + weight
                
                if distance < distances[neighbor]:
                    distances[neighbor]=distance
                    heapq.heappush(queue,(distance,neighbor))


    return distances

N,M,K,X = map(int,input().split())

graph = [[0]*N for _ in range(N)]


for i in range(M):
    x, y = map(int,input().split())
    graph[x-1][y-1] = 1

isNone = 1
distances = dijkstra(graph,X)
print(distances)
print(graph)
if i in range(1,N+1):
    if distances[i] == K:
        print(i)
        isNone=0
if isNone:
    print(-1)

 

이 문제는 2D array를 통해서 문제를 풀었다. (연결리스트로 하면 시간복잡도가 낮아지니 이후에는 연결리스트를 이용했다)

 

우선 우리는 다익스트라 알고리즘을 이용해서 X에서 시작하는 최단거리 리스트를 만들어야 한다.

아래는 다익스트라 알고리즘의 형태이다. 나는 함수 안에 distances 리스트를 만들었는데 이는 함수 밖에서 선언해도 무관하다.

def dijkstra(graph, start):
    distances=[float('inf')]*len(graph)
    distances[start-1]=0

    queue = [(0,start)]

    while queue :
        curr_distance, curr_node = heapq.heappop(queue)
        if curr_distance > distances[curr_node] :
            continue

        for neighbor in range(len(graph)):
            weight = graph[start][neighbor]
            
            if weight != 0 :
                distance = curr_distance + weight
                
                if distance < distances[neighbor]:
                    distances[neighbor]=distance
                    heapq.heappush(queue,(distance,neighbor))

 

 

distances를 무한한 값으로 초기 화 하고 첫번째 경로를 queue에 담는다. 

 -> 다익스트라 알고리즘은 queue가 비워져 있을경우 break를 하기 때문에 while문 이전에 queue를 넣어주어야 한다.

 

def dijkstra(graph, start):
    distances=[float('inf')]*len(graph)
    distances[start-1]=0

    queue = [(0,start)]

 

큐가 채워져 있을 경우 큐에서 heappop을 해서 현재 경로와 노드를 갱신한다.

이때 현재 경로가 distances리스트 값보다 더 클 경우 continue한다.

우리는 큐를 돌면서 계속 distances를 최신화 함으로 더 비용이 큰값은 들어가지 않는다.

    while queue :
        curr_distance, curr_node = heapq.heappop(queue)
        if curr_distance > distances[curr_node] :
            continue

 

다익스트라 알고리즘은 이웃노드를 확인해서 이웃노드의 가중치가 기존에 업데이트한 가중치보다 작을경우 업데이트를 하고 해당 노드를 우선순위 큐에 넣어서 이를 반복한다.

우선순위 큐에 넣을때는 항상 거리를 먼저 넣어야 한다. 큐는 먼저 넣은 거리를 기준으로 우선순위 큐를 만들기 때문이다.

        for neighbor in range(len(graph)):
            weight = graph[start][neighbor]
            
            if weight != 0 :
                distance = curr_distance + weight
                
                if distance < distances[neighbor]:
                    distances[neighbor]=distance
                    heapq.heappush(queue,(distance,neighbor))

 

위에서 다익스트라 알고리즘을 실행했다면 우리는 거리가 K인 노드를 출력하고 싶음으로 K를 찾는 반복문을 돌린다.

만약 반복문이 다 돌았는데도 못찾았을 경우 isNone이라는 변수를 통해서 -1을 출력할 수 있도록 한다.

isNone = 1
distances = dijkstra(graph,X)
print(distances)
print(graph)
for i in range(1,N+1):
    if distances[i] == K:
        print(i)
        isNone=0
if isNone:
    print(-1)