백준 1753

시간 제한메모리 제한제출정답맞힌 사람정답 비율
1 초256 MB171625507652558025.010%

문제

방향그래프가 주어지면 주어진 시작점에서 다른 모든 정점으로의 최단 경로를 구하는 프로그램을 작성하시오. 단, 모든 간선의 가중치는 10 이하의 자연수이다.

입력

첫째 줄에 정점의 개수 V와 간선의 개수 E가 주어진다. (1 ≤ V ≤ 20,000, 1 ≤ E ≤ 300,000) 모든 정점에는 1부터 V까지 번호가 매겨져 있다고 가정한다. 둘째 줄에는 시작 정점의 번호 K(1 ≤ K ≤ V)가 주어진다. 셋째 줄부터 E개의 줄에 걸쳐 각 간선을 나타내는 세 개의 정수 (u, v, w)가 순서대로 주어진다. 이는 u에서 v로 가는 가중치 w인 간선이 존재한다는 뜻이다. u와 v는 서로 다르며 w는 10 이하의 자연수이다. 서로 다른 두 정점 사이에 여러 개의 간선이 존재할 수도 있음에 유의한다.

출력

첫째 줄부터 V개의 줄에 걸쳐, i번째 줄에 i번 정점으로의 최단 경로의 경로값을 출력한다. 시작점 자신은 0으로 출력하고, 경로가 존재하지 않는 경우에는 INF를 출력하면 된다.


답안

import math
from collections import defaultdict
import heapq
import sys

def dijkstra(graph, start, vertex_count):
    distances = [math.inf] * (vertex_count+1)
    distances[start] = 0
    visited = [False] * (vertex_count+1)
    pqueue = []
    heapq.heappush(pqueue, [0, start])

    while pqueue:
        _, src = heapq.heappop(pqueue)
        if visited[src]:
            continue
        visited[src] = True
        for dest in graph[src].keys():
            weight = graph[src][dest]
            if distances[src] + weight < distances[dest]:
                distances[dest] = distances[src] + weight
                heapq.heappush(pqueue, [distances[dest], dest])

    return distances[1:]

vertex_count, edge_count = map(int, sys.stdin.readline().split())
graph = defaultdict(lambda: defaultdict(lambda: math.inf))
start = int(sys.stdin.readline())
for i in range(edge_count):
    src, dest, weight = map(int, sys.stdin.readline().split())
    graph[src][dest] = min(graph[src][dest], weight)

for distance in dijkstra(graph, start, vertex_count):
    print('INF') if distance is math.inf else print(distance)

풀이

이 풀이는 다익스트라 알고리즘을 이용한 풀이이다.

다익스트라 알고리즘을 간단하게 설명하면, BFS와 그리디로 구성된 최단 경로 탐색 알고리즘으로, 한 정점으로부터 다른 모든 정점까지의 최단 경로를 구할 수 있다.

다익스트라 알고리즘은 다음과 같은 과정을 거친다.

  1. 출발 노드를 설정하고, 출발 노드로부터 각 노드까지의 거리를 저장한다.
  2. 방문하지 않은 노드들 중 가장 거리가 짧은 노드를 선택해 방문한다.
  3. 해당 노드를 거쳐 각 노드까지의 거리가 짧아진다면 갱신한다.
  4. 모든 노드를 방문할 때까지 2번과 3번을 반복한다.
vertex_count, edge_count = map(int, sys.stdin.readline().split())
graph = defaultdict(lambda: defaultdict(lambda: math.inf))
start = int(sys.stdin.readline())
for i in range(edge_count):
    src, dest, weight = map(int, sys.stdin.readline().split())
    graph[src][dest] = min(graph[src][dest], weight)

위와 같이 필요한 변수들(vertex_count, edge_count, graph, start)을 초기화해준다.

각각은 정점의 개수, 간선의 개수, 그래프를 나타내는 2차원 인접 배열, 시작 정점을 의미한다.

def dijkstra(graph, start, vertex_count):
    distances = [math.inf] * (vertex_count+1)
    distances[start] = 0
    visited = [False] * (vertex_count+1)
    pqueue = []
    heapq.heappush(pqueue, [0, start])
		...

distances 리스트는 시작 정점으로부터 k번째 정점으로 가는 최소 거리를 나타낸다. 처음에는 무한대(math.inf)로 초기화해준다.

visited 리스트를 통해 각 노드를 방문하였는지 여부를 저장한다.

pqueue라는 이름의 priority queue를 통해, 방문하지 않은 노드들 중 가장 거리가 짧은 노드를 선택하도록 한다.

	
		...
    while pqueue:
        _, src = heapq.heappop(pqueue)
        if visited[src]:
            continue
        visited[src] = True
        for dest in graph[src].keys():
            weight = graph[src][dest]
            if distances[src] + weight < distances[dest]:
                distances[dest] = distances[src] + weight
                heapq.heappush(pqueue, [distances[dest], dest])
		...   

모든 노드를 방문할 때까지 노드 방문과 거리 갱신을 반복한다.

만약 현재 노드를 거쳐 특정 노드로 가는 거리(distances[src] + weight)가 현재까지 밝혀진 특정 노드까지의 최단 거리보다 더 작을 경우, 거리를 갱신한다.

def dijkstra(...):
		...
		return distances[1:]

...

for distance in dijkstra(graph, start, vertex_count):
    print('INF') if distance is math.inf else print(distance)

시작 노드로부터 1~n번 노드까지의 최단거리를 나타낸 배열(distances[1:])을 이용해, 만약 도달 불가능한 경우(방문하지 못한 경우) ‘INF’를, 그렇지 않은 경우 최단 경로를 출력한다.

profile
이예찬

0개의 댓글