Koala - 9기/코딩테스트 준비 스터디
[백준/Python] 17396번 백도어
브로리카
2023. 2. 16. 13:45
다익스트라를 사용해 푸는 문제
평범한 골드 5 문제라고 생각했다가 시간 초과 2번 맞고 다익스트라를 잘못 이해하고 있었단 걸 알았다.
시간 초과의 원인
https://www.acmicpc.net/board/view/34516
7번 항목을 고려하지 않고 코드를 작성했기 때문에 시간 초과를 받았다.
평소에는 저 부분을 고려하지 않고 풀어도 잘 통과했다.
하지만, 정점이 10만개가 되니, 불필요하게 탐색되는 정점의 수가 누적돼 시간 초과로 이어졌다.
어떻게 해결했나?
매우 간단하다.
기존 코드에 2줄만 추가하면 됐다.
문제의 코드
while q:
hereC, here = heapq.heappop(q)
for there, thereC in graph[here]:
cost = thereC + hereC
if dist[there] > cost:
dist[there] = cost
heapq.heappush(q, (cost, there))
개선한 코드 (3~4번 라인)
while q:
hereC, here = heapq.heappop(q)
if dist[here] < hereC: # 추가된 부분
continue
for there, thereC in graph[here]:
cost = thereC + hereC
if dist[there] > cost:
dist[there] = cost
heapq.heappush(q, (cost, there))
전체 코드
import sys, heapq
input = sys.stdin.readline
def main():
N, M = map(int, input().split())
graph = [[] for _ in range(N)]
wards = list(map(int, input().split()))
wards[-1] = 0
for _ in range(M):
u, v, c = map(int, input().split())
if wards[u] or wards[v]:
continue
graph[u].append((v, c))
graph[v].append((u, c))
dist = [float("inf")] * N
dist[0] = 0
q = [(0, 0)]
while q:
hereC, here = heapq.heappop(q)
if dist[here] < hereC:
continue
for there, thereC in graph[here]:
cost = thereC + hereC
if dist[there] > cost:
dist[there] = cost
heapq.heappush(q, (cost, there))
if dist[-1] == float("inf"):
print(-1)
else:
print(dist[-1])
if __name__ == "__main__":
main()