늑대가 현재 빠르게 이동하는지, 느리게 이동하는지에 따라서 cost<<2인지, cost인지 결정.
다익스트라를 따로 돌려주기.
늑대의 경우엔 현재 빠르게 달릴 수 있는지, 없는지 fast라는 state를 힙에 같이 넣는다.
fast이면 간선 비용이 절반인 cost를 가져오고, fast가 아니면 4배가 된 cost를 가져온다.
전체코드
파이썬
import sys
import heapq as hq
input = lambda : sys.stdin.readline().strip()
INF = float('inf')
# 간선 입력받기
N,M = map(int,input().split())
G_wolf = [[[] for _ in range(N+1)] for _ in range(2)]
G_fox = [[] for _ in range(N+1)]
for _ in range(M):
a,b,cost = map(int,input().split())
for x,y in [(a,b),(b,a)]:
G_wolf[1][x].append((y,cost))
G_wolf[0][x].append((y,cost<<2))
G_fox[x].append((y,cost<<1))
# 늑대 다익스트라
V_wolf = [[INF]*(N+1) for _ in range(2)]
V_wolf[0][1] = 0
PQ_wolf = [(0,1,1)]
while PQ_wolf:
t,x,fast = hq.heappop(PQ_wolf)
# x이고 그 노드에 방문했을때는 fast^1 상태였다.
if t > V_wolf[fast^1][x]:
continue
# fast 상태에 따른 간선 선택만 잘 해주기
for nx,cost in G_wolf[fast][x]:
nt = t + cost
if nt < V_wolf[fast][nx]:
V_wolf[fast][nx] = nt
hq.heappush(PQ_wolf,(nt,nx,fast^1))
# 여우 다익스트라 (기존이랑 동일)
V_fox = [INF]*(N+1)
V_fox[1] = 0
PQ_fox = [(0,1)]
while PQ_fox:
t,x = hq.heappop(PQ_fox)
if t > V_fox[x]:
continue
for nx,cost in G_fox[x]:
nt = t + cost
if nt < V_fox[nx]:
V_fox[nx] = nt
hq.heappush(PQ_fox,(nt,nx))
answer = 0
for i in range(1,N+1):
if V_fox[i] < min(V_wolf[0][i], V_wolf[1][i]):
answer += 1
print(answer)
코멘트
간선이 많아서 더미 데이터가 힙에 자꾸 쌓인다.
t,x,fast 정보를 가진게 힙에서 나왔을 때, 이 정보는 V[fast^1][x]에 방문했기 때문에 토글을 해줘야한다.
이 부분을 캐치를 못해서 오류나길래 빼고 돌렸는데 TLE 나서 조금 헤멨다.
어려운 다익문제는 아닌데 생각해볼 거리가 많은 문제.
간선 cost같은 경우, 기본적으로 cost<<1로 만들어놓으면 짝수라서, >>1로 땡겨올 때도 큰 제약 없이 풀이가 가능하다.
자바로 풀 때 오버플로우 생각하기.
사실 입력 범위보면 int 안써도 될거같긴 함.
PQ에 들어가는 t에서 혼자 long쓰면 클래스 만들고 뻘짓해야함.
간선 비용 최대 4배로 한 G_wolf에서 worst로 잡아서 생각했을 때, 40만 * 4천이라 16억이라 long 안쓰고 int에서 가능하다.
댓글