[프로그래머스] 부대복귀

문제

• 지역은 번호로 구분되고 두  지역 간 길을 통과하는 데 1시간이 걸린다. roads 변수에는 길이 있는 두 지역의 번호를 포함하고 있다.
  • 3 <= 지역의 수 <= 100,000
  • 2 <= 길의 수 <= 500,000
• 강철부대의 지역이 주어졌을 때 도착지까지 복귀할 수 있는 최단시간 찾기
• 복귀할 수 없다면 -1
  
  

 

풀이

• destination 에서의 각 최단거리 이므로 다익스트라 알고리즘 사용하기
• 시간복잡도 : O(ElogV) 이므로 가능
  V = 100,000 / E = 500,000

 

import heapq
import sys
input = sys.stdin.readline
INF = int(1e9)

def solution(n, roads, sources, destination):
    distance = [INF] * (n+1)

    # 그래프 만들기
    graph = [[] for _ in range(n+1)]
    for a, b in roads :
        graph[a].append(b)
        graph[b].append(a)

    # dijkstra
    q = []
    heapq.heappush(q, (0, destination))
    distance[destination] = 0
    while q:
        dist, now = heapq.heappop(q)

        if distance[now] < dist :
            continue

        for i in graph[now]:
            if dist+1 < distance[i] :
                distance[i] = dist+1
                heapq.heappush(q, (dist+1, i))

    answer = []
    # 복귀 시간 찾기
    for s in sources :
        if distance[s] == INF :
            answer.append(-1)
        else:
            answer.append(distance[s])

    return answer

 

 

풀이 2

• BFS로 최단거리를 구하되, 한 번에 모든 지점까지의 최단거리 구하기
• 다익스트라 방법보다 빠르다

 

from collections import deque

def solution(n, roads, sources, destination):
    graph = [[] for _ in range(n+1)]
    for a, b in roads:
        graph[a].append(b)
        graph[b].append(a)
           
    #최단거리 찾기
    #도착지부터 모든 지점의 최단거리 구함 (첫 시작=destination) 
    queue = deque([destination])
    
    visit = [-1]*(n+1)
    visit[destination] = 0
    
    #BFS 한번에 구하기
    while queue:
        cur = queue.popleft()
        
        for n in graph[cur]:
            if visit[n] == -1:
                visit[n] = visit[cur] + 1
                queue.append(n)
                
    return [visit[i] for i in sources]

 

 

References

• 프로그래머스 Lv3 부대복귀
• 0sun-creator