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

Problem : https://school.programmers.co.kr/learn/courses/30/lessons/132266

Status : Solved 

Time : 00:14:44

 

---

 

문제 설명

 

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.

강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources, 강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.

 

입력 및 출력

 

제한사항

• 3 ≤ n ≤ 100,000
  • 각 지역은 정수 1부터 n까지의 번호로 구분됩니다.
• 2 ≤ roads의 길이 ≤ 500,000
  • roads의 원소의 길이 = 2
  • roads의 원소는 [a, b] 형태로 두 지역 a, b가 서로 왕복할 수 있음을 의미합니다.(1 ≤ a, b ≤ n, a ≠ b)
  • 동일한 정보가 중복해서 주어지지 않습니다.
    • 동일한 [a, b]가 중복해서 주어지지 않습니다.
    • [a, b]가 있다면 [b, a]는 주어지지 않습니다.
• 1 ≤ sources의 길이 ≤ 500
  • 1 ≤ sources[i] ≤ n
• 1 ≤ destination ≤ n

 

입출력 예

n	roads	sources	destination
3	[[1, 2],[2,3]]	[2,3]	1
5	[[1,2], [1,4],[2,4],[2,5],[4,5]]	[1,3,5]	5

 

 

---

 

풀이

 

경로 찾기 알고리즘 문제이다. 출발지가 리스트 단위로 주어지고, 도착지가 정수, 즉 하나로 주어지는 형태이다. 따라서 모든 출발지와 도착지의 거리관계를 구하는 플로이드-와셜 알고리즘을 사용해봄직 하지만, 당연히 시간초과가 나게 된다.

 

따라서 우리는 출발지와 도착지를 바꿔서 볼 필요가 있다. 도착지 관점에서 출발지 리스트를 바라보면 이는 간단히 다익스트라 알고리즘으로  구현할 수 있다. 다익스트라 알고리즘으로 도착지에서 각 노드까지의 최단거리를 구해나가며 그 리스트를 저장하고, 추후 출발지에 따른 거리를 반환하면 되는 문제이다.

 

풀이 코드

from heapq import heappush,heappop

MAX = float('inf')
def solution(n, roads, sources, destination):
    map_dict = { key : [] for key in range(1, n+1)}
    
    for a, b in roads :
        map_dict[a].append(b)
        map_dict[b].append(a)
    
    dist_list = [MAX]*(n+1)
    q = [(0, destination)]
    
    while q :
        dist, node = heappop(q)
        if dist_list[node] <= dist :
            continue
        dist_list[node] = dist
        for next_node in map_dict[node] :
            heappush(q, (dist+1, next_node))
    
    answer_list = []
    for source in sources :
        if source == destination :
            answer_list.append(0)
        elif dist_list[source] == MAX :
            answer_list.append(-1)
        else :
            answer_list.append(dist_list[source])
                
    return answer_list

풀이 완료!