[백준 16118번] 달빛 여우 (Python3)

import sys
input = sys.stdin.readline
from heapq import *

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

graph = [[] for i in range(N)]
for _ in range(M):
  a,b,c = map(int,input().split())
  graph[a-1].append((b-1,c*2))
  graph[b-1].append((a-1,c*2))

fox = [1e12]*N; hq = [(0,0)]
while hq:
  w,now = heappop(hq)
  if fox[now] < w:
    continue
  for next,w1 in graph[now]:
    if fox[next] > w+w1:
      fox[next] = w+w1
      heappush(hq,(w+w1,next))

wolf = [[1e12]*N for i in range(2)]; hq = [(0,0,0)]
while hq:
  w,now,s = heappop(hq)
  if wolf[s][now] < w:
    continue
  for next,w1 in graph[now]:
    if s:
      w1 <<= 1
    else:
      w1 >>= 1
    if wolf[s^1][next] > w+w1:
      wolf[s^1][next] = w+w1
      heappush(hq,(w+w1,next,s^1))

print(sum([int(fox[i]<min(wolf[0][i],wolf[1][i])) for i in range(1,N)]))

간단한 다 익스트라 문제

...이지만 파이썬의 경우 까다로울 수 있는 문제이다.

 

추가시간이 없는 탓에 파이썬을 사용하면 억까당할 수 있다. 풀이방식이 똑같아도 사소한 차이로 시간초과가 발생하기 때문이다. 

나의 경우에는 다 익스트라 구현 방식때문에 시간초과가 발생하였다.

시간초과가 발생한 코드

통과된 코드

 

풀이를 설명하면,

1. 그래프를 입력받을 때 가중치에 *2를 해준다.

2. 여우의 각 노드에 대한 최단거리를 다 익스트라로 구한다.

3. 늑대의 각 노드에 대한 최단거리를 다 익스트라로 구한다. 이때 DP 배열을 N*2의 2차원으로 만든다.

DP[0][now]는 다음번에 두배속으로 이동하는 경우의 최단거리, DP[1][now는 다음번에 절반의 속도로 이동하는 경우의 최단거리를 의미한다.

4. 여우가 더 빠르게 이동할 수 있는 노드의 개수를 출력한다.

 

 

문제 자체는 매우 간단했지만, 내 다 익스트라 구현방식에 문제가 있었다는 사실을 알게해준 좋은 문제였다. 그동안 다 익스트라 문제를 몇십개를 풀었는데 이제서야 깨닫다니....

 

오늘의 교훈) 다 익스트라 구현 방식을 고민해보자