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()