[Python/파이썬] 백준 16118번 달빛 여우

16118번: 달빛 여우

 

문제

관악산 기슭에는 보름달을 기다리는 달빛 여우가 한 마리 살고 있다. 달빛 여우가 보름달의 달빛을 받으면 아름다운 구미호로 변신할 수 있다. 하지만 보름달을 기다리는 건 달빛 여우뿐만이 아니다. 달빛을 받아서 멋진 늑대인간이 되고 싶어 하는 달빛 늑대도 한 마리 살고 있다.

관악산에는 1번부터 N번까지의 번호가 붙은 N개의 나무 그루터기가 있고, 그루터기들 사이에는 M개의 오솔길이 나 있다. 오솔길은 어떤 방향으로든 지나갈 수 있으며, 어떤 두 그루터기 사이에 두 개 이상의 오솔길이 나 있는 경우는 없다. 달빛 여우와 달빛 늑대는 1번 나무 그루터기에서 살고 있다.

보름달이 뜨면 나무 그루터기들 중 하나가 달빛을 받아 밝게 빛나게 된다. 그러면 달빛 여우와 달빛 늑대는 먼저 달빛을 독차지하기 위해 최대한 빨리 오솔길을 따라서 그 그루터기로 달려가야 한다. 이때 달빛 여우는 늘 일정한 속도로 달려가는 반면, 달빛 늑대는 달빛 여우보다 더 빠르게 달릴 수 있지만 체력이 부족하기 때문에 다른 전략을 사용한다. 달빛 늑대는 출발할 때 오솔길 하나를 달빛 여우의 두 배의 속도로 달려가고, 그 뒤로는 오솔길 하나를 달빛 여우의 절반의 속도로 걸어가며 체력을 회복하고 나서 다음 오솔길을 다시 달빛 여우의 두 배의 속도로 달려가는 것을 반복한다. 달빛 여우와 달빛 늑대는 각자 가장 빠르게 달빛이 비치는 그루터기까지 다다를 수 있는 경로로 이동한다. 따라서 둘의 이동 경로가 서로 다를 수도 있다.

출제자는 관악산의 모든 동물을 사랑하지만, 이번에는 달빛 여우를 조금 더 사랑해 주기로 했다. 그래서 달빛 여우가 달빛 늑대보다 먼저 도착할 수 있는 그루터기에 달빛을 비춰 주려고 한다. 이런 그루터기가 몇 개나 있는지 알아보자.

입력

첫 줄에 나무 그루터기의 개수와 오솔길의 개수를 의미하는 정수 N, M(2 ≤ N ≤ 4,000, 1 ≤ M ≤ 100,000)이 주어진다.

두 번째 줄부터 M개의 줄에 걸쳐 각 줄에 세 개의 정수 a, b, d(1 ≤ a, b ≤ N, a ≠ b, 1 ≤ d ≤ 100,000)가 주어진다. 이는 a번 그루터기와 b번 그루터기 사이에 길이가 d인 오솔길이 나 있음을 의미한다.

출력

첫 줄에 달빛 여우가 달빛 늑대보다 먼저 도착할 수 있는 나무 그루터기의 개수를 출력한다.

코드

from heapq import heappop, heappush
import sys
input = sys.stdin.readline

n, m = map(int, input().split())
road = [[] for _ in range(n+1)]
for _ in range(m):
    a, b, d = map(int, input().split())
    road[a].append((b, d*2))
    road[b].append((a, d*2))

# 여우
fox_dist = [int(1e9)] * (n+1)
hq = [(0, 1)]
while hq:
    now_d, now = heappop(hq)
    if fox_dist[now] < now_d:
        continue
    for nx, nx_d in road[now]:
        cost = now_d + nx_d
        if fox_dist[nx] > cost:  # now 거쳐서 nx 가는게 더 빠른 경우
            fox_dist[nx] = cost
            heappush(hq, (cost, nx))

# 늑대
wolf_dist = [[int(1e9)]*2 for _ in range(n+1)]
hq = [(0, 1, True)]
while hq:
    now_d, now, fast = heappop(hq)
    # now->nx로 2배 속도로 움직임
    if fast:
        if wolf_dist[now][1] < now_d:
            continue
        for nx, nx_d in road[now]:
            cost = now_d + nx_d // 2
            if wolf_dist[nx][0] > cost:
                wolf_dist[nx][0] = cost
                heappush(hq, (cost, nx, False))
    # now->nx로 0.5배 속도로 움직임
    else:
        if wolf_dist[now][0] < now_d:
            continue
        for nx, nx_d in road[now]:
            cost = now_d + nx_d * 2
            if wolf_dist[nx][1] > cost:
                wolf_dist[nx][1] = cost
                heappush(hq, (cost, nx, True))

ans = 0
for i in range(2, n+1):
    if fox_dist[i] < wolf_dist[i][0] and fox_dist[i] < wolf_dist[i][1]:
        ans += 1
print(ans)

나중에 늑대의 속도를 계산할 때 2로 나누기도 하는데 이 때 나누기 연산 후에도 정수일 수 있도록 그루터기 간의 경로를 입력 받을 때 X2를 해준다.

한 그루터기에서 다른 모든 그루터기들까지의 최단 거리를 구하면 되는 문제이기 때문에 다익스트라 알고리즘을 사용한다. 여우는 일반적인 다익스트라 알고리즘을 사용해서 구해주면 된다. 늑대는 이번에 2배의 속도로 움직일지 0.5배의 속도로 움직일지 2가지 경우로 나눠서 봐야 한다. 그렇기 때문에 최단 거리를 저장해줄 때도 2가지 경우로 나누어서 저장해주어야 해서 배열을 N행 2열로 만들었다.

 

알고리즘 자체는 어렵지 않은데 시간 초과 때문에 고생했던 문제였다. 질문 게시판도 참고하고 구글링도 해서 겨우 시간 초과 해결했다...