본문 바로가기
나동빈 with 파이썬/실전문제 풀이

[Ch. 9 - 최단 경로] 전보

by yewonnie 2022. 2. 17.

My code

import heapq
INF = int(1e9)

n, m, c = map(int,input().split())

graph = [[] for i in range(n + 1)]
distance = [INF] * (n + 1)

for i in range(m):
    x, y, z = map(int,input().split())
    graph[x].append((y, z))

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

dijkstra(c)

count = 0
max_dist = 0
for i in range(2, n + 1):
    if distance[i] != INF:
        count += 1
        max_dist = max(distance[i], max_dist)

print(count, max_dist)

Answer

import heapq
INF = int(1e9)

n, m, start = map(int,input().split())

graph = [[] for i in range(n + 1)]
distance = [INF] * (n + 1)

for i in range(m):
    x, y, z = map(int,input().split())
    graph[x].append((y, z))

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

dijkstra(start)

count = 0
max_distance = 0
for d in distance:
    if d != INF:
        count += 1
        max_distance = max(max_distance, d)

print(count - 1, max_distance)

Code Review

이번 문제는 C 라는 도시에서 최대한 많은 도시로 메시지를 보내고자 할 때 메시지를 받게 되는 도시의 개수와 도시들이 모두 메시지를 받는 데까지 걸리는 시간은 얼마인지 구해보는 문제였습니다. C 에서 출발하여 나머지 도시까지의 최단 경로를 구하는 문제이므로 Dijkstra 알고리즘으로 이 문제를 해결 할 수 있었습니다. 
따라서 Dijkstra 알고리즘을 이용해 먼저 C 에서 다른 도시들까지의 최단경로를 구해 1차원 리스트 distance 에 저장해주었습니다. 
처음 문제를 읽었을 때 도시들이 모두 메시지를 받는 데까지 걸리는 시간을 구하라는 뜻을 잘못이해했었습니다. 따라서 각각의 도시들이 메시지를 받는데 걸리는 시간 즉, distance 의 값이 INF 가 아니라면 그 안의 값을 모두 더해주었었습니다. 하지만 출력 예시를 보니 도시들이 모두 메시지를 받는 데까지 걸리는 시간이 이 뜻이 아니라 메시지를 받은 도시들 중 최대로 걸린 시간이 몇 시간이 구하는 것이었습니다. 왜냐하면 그 시간 안에 다른 도시들도 메시지를 받았을 것이기 때문입니다. 즉, distance 의 값 중 최대값을 찾으면 되는 것이었습니다. 
따라서 distance 를 2번째 부터 탐색하여 (시작 노드는 제외해야 하므로) 만약 값이 INF 가 아니라면 count 를 1 증가시켜주었습니다. 또한 max 함수를 이용하여 distance 의 값중 가장 큰 값을 찾아주어 count 의 값과 함께 출력해주었습니다. 

댓글