[이코테] 8.1 최단 경로 실전문제

미래 도시


해당 문제는 플로이드 워셜 최단 경로 알고리즘의 대표적인 문제이다.

우선 입력으로 주어진 노드의 개수가 100 이하, 간선의 개수가 100 이하이므로 O(N^3)알고리즘을 사용할 수 있다.

문제는 start → K → X 이므로 start → K 에 대한 최단 경로와 K → X에 대한 최단 경로 두가지가 필요하다. Dijkstra와 같은 single to all 을 2번 사용하여 구할 수 있지만 O(N^3)알고리즘을 사용해도 시간에 문제가 없으므로 모든 노드에서 모든 노드로의 최단 경로를 구해주는 플로이드-워셜 알고리즘을 사용하면 간단하니까 유리하다.

import sys

def floyd_warshall(n):
    for k in range(1, n + 1):
        for i in range(1, n + 1):
            for j in range(1, n + 1):
                graph[i][j] = min(graph[i][j], graph[i][k] + graph[k][j])

def solution(n, x, k):
    floyd_warshall(n)
    answer = graph[1][k] + graph[k][x]
    if answer >= INF:
        return -1
    return answer

if __name__ == "__main__":
    N, M = map(int, sys.stdin.readline().rstrip().split())
    INF = int(1e9)
    graph = [[INF] * (N+1) for _ in range(N+1)]

    for a in range(1, N + 1):
        for b in range(1, N + 1):
            if a == b:
                graph[a][b] = 0

    for _ in range(M):
        s, e = map(int, sys.stdin.readline().rstrip().split())
        graph[s][e] = 1
        graph[e][s] = 1

    X, K = map(int, sys.stdin.readline().rstrip().split())

    print(solution(N, X, K))

전보


해당 문제는 다익스트라의 대표적인 유형이다.

우선 입력을 보면 노드의 개수는 30000개 이하, 간선의 개수는 20만개 이하이다. 플로이드 워셜 알고리즘을 사용하면 시간 초과가 발생할 것이다.

우선 문제에서 “C라는 도시에서 최대한 많은 도시로 메세지를 보내고자 한다” 라는 부분에서 Single Source → All Destination이라는 힌트를 얻을 수 있다. 따라서 이에 해당하는 다익스트라를 떠올릴 수 있다.

앞서 본 입력에 대해 2가지 유형의 다익스트라중, 노드의 개수가 10000개가 넘기 때문에 최소 힙을 사용하는 개선된 다익스트라 알고리즘을 사용하여 구현해야 한다.

import sys
import heapq

def dijkstra(start):
    heap = []
    heapq.heappush(heap, (0, start))
    distance[start] = 0
    while heap:
        dist, now = heapq.heappop(heap)
        if distance[now] < dist:
            continue
        for a in graph[now]:
            cost = dist + a[1]
            if cost < distance[a[0]]:
                distance[a[0]] = cost
                heapq.heappush(heap, (cost, a[0]))

def solution(start):
    dijkstra(start)
    _max_cnt = len(distance)
    _max_time = 0
    for d in distance:
        if d == INF:
            _max_cnt -= 1
        elif d > _max_time:
            _max_time = d
    print(_max_cnt-1, _max_time)

if __name__ == "__main__":
    N, M, C = map(int, sys.stdin.readline().rstrip().split())
    INF = 1001
    graph = [[] * (N + 1) for _ in range(N + 1)]
    distance = [INF] * (N + 1)
    for _ in range(M):
        X, Y, Z = map(int, sys.stdin.readline().rstrip().split())
        graph[X].append((Y, Z))
    solution(C)