[python] algospot FIRETRUCK

문제

-> 문제 링크 : https://algospot.com/judge/problem/read/FIRETRUCKS

algospot.com :: FIRETRUCKS

 

문제풀이

이 문제는 소방서에서 불난 지점까지의 최단 거리들의 합을 구하는 문제이다.

소방서의 개수가 최대 1000개이기 때문에 모든 소방서에서 최단경로 알고리즘을 돌린 뒤 최단 거리를 찾으면 시간 초과가 발생한다.

다익스트라 알고리즘을 한 번만 돌려 풀 수 있는 방법을 생각해야 한다.

 

이 문제의 핵심은 모든 소방서들은 최단 거리가 0이라는 것이다. 

한 소방서에서 시작하는 다익스트라 알고리즘을 돌려 경로의 가중치가 쌓이더라도 다른 소방서에 도착하면 최단 거리는 0이 된다.

 

기본 다익스트라와 같이 현재 정점까지의 최단 경로와 다음 정점으로 가는 cost를 더한 값이 다음 정점의 최단 경로보다 작다면 우선순위 큐에 넣어준다. 하지만 이때 다음 정점이 소방서라면 무조건 다음 정점의 최단경로를 0으로 세팅하고 우선순위 큐에 넣어준다.

이렇게 하면 자동적으로 이미 최단 거리로 계산된 경로들도 이 소방서를 기준으로 다시 최단 경로를 계산하게 된다. 

 

코드

import sys
from collections import defaultdict
import heapq

T = int(sys.stdin.readline().rstrip())

def dijkstra(fires, stations) :
    pq = []
    start = stations[0]
    heapq.heappush(pq, (start,0))

    dist = [float('inf')] * (V+1)
    dist[start] = 0

    isStation = [False] * (V+1)
    for station in stations :
        isStation[station] = True

    isFire = [False] * (V+1)
    for fire in fires :
        isFire[fire] = True

    while pq :
        here, cost = heapq.heappop(pq)

        if dist[here] < cost :
            continue

        for next, next_cost in graph[here] :
            if isStation[next] :
                if dist[next] != 0 :
                    dist[next] = 0
                    heapq.heappush(pq, (next, 0))
                    
            else :
                next_dist = next_cost + dist[here]
                if dist[next] > next_dist:
                    dist[next] = next_dist
                    heapq.heappush(pq, (next, next_dist))

    return dist


for _ in range(T) :
    V, E, N, M = map(int, sys.stdin.readline().rstrip().split(" "))

    graph = defaultdict(list)
    for _ in range(E) :
        a,b,t = map(int, sys.stdin.readline().rstrip().split(" "))
        graph[a].append((b, t))
        graph[b].append((a, t))

    fires = list(map(int,sys.stdin.readline().rstrip().split()))
    stations = list(map(int,sys.stdin.readline().rstrip().split()))

    distances = dijkstra(fires, stations)
    answer = 0
    for fire in fires :
        answer += distances[fire]

    print(answer)

 

 

코드에 잘못된 부분이나 개선 사항을 댓글로 알려주시면 감사하겠습니다.