부대 복귀

링크

코딩테스트 연습 - 부대복귀

문제에 대한 이해

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.
강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources, 강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.

• 각 위치에서 강철부대로 복귀할 수 있는 최단 시간을 구한다.
• 길이 없는 경우도 존재한다. 이럴 때는 -1을 출력한다.
• 주어지는 roads는 [a, b] 형태로 주어지며 항상 왕복 가능하다는 것을 뜻한다.

문제의 예시를 보면 알 수 있지만, 적군의 방해로 인해 되돌아오는 경로가 없어져 라는 것에 큰 의미를 둘 필요는 없다. 단순히 이는 destination으로 이어지는 길이 없을 수 있다는 뜻으로 해석할 수 있다.

접근 방법

강철 부대와 다른 지점 간 연결 정보가 필요하다.

연결 정보가 필요한 이유는, 강철 부대로부터 다른 노드를 방문하기 위해서는 서로 연결되어 있어야 하며, 이를 통해 두 지점 사이의 거리를 알 수 있기 때문이다.

문제의 input 중, 2차원 정수 배열을 roads로 담아서 주고, 이는 두 지점 간에 왕복이 가능하다는 뜻이므로, 이 배열을 순회하면서 인접 리스트를 완성해준다.

Python
def solution(n, roads, sources, destination):
    # 인접 리스트 만들기
    adj = [[] for _ in range(n + 1)]

    for a, b in roads:
        # 서로 왕복할 수 있어야 하므로 각 인덱스 모두 목적지를 넣어준다.
        adj[a].append(b)
        adj[b].append(a)

각 노드와 다른 노드 사이의 연결 정보를 얻었으면, 해당 정보를 이용하여 강철 부대의 위치로부터 다른 지점까지의 거리를 구한다.

복귀하는 경로를 구하는 것은 곧 강철부대로부터 해당 지점까지의 거리와 동일하다는 뜻이기 때문에, 그리고 모든 지점에서 강철 부대로 복귀하는 거리를 구하는 것이 문제이므로, 이는 강철 부대를 시작점으로 하고 해당 지점으로부터 각 지점까지의 거리를 구하라는 뜻과 같다.

거리를 구하기 위하여 크게 3가지 방법을 사용할 수 있다.

1. dfs
2. bfs
3. dijkstra(heap 사용)

우선 dfs와 bfs는 두 노드 사이의 간선의 가중치가 1인 경우에만 이용할 수 있는 방법이다. 이 문제에서는 두 지역간의 길을 통과하는데 걸리는 시간은 모두 1로 동일하기 때문에, 두 방법 모두 이용할 수 있다.

다만, 문제에서 노드의 개수가 최대 10만개이기 때문에, 재귀적 방법인 dfs를 이용한다면 파이썬의 경우 stackoverflow error가 발생할 수 있다. 파이썬의 허용된 재귀 깊이는 1000이기 때문이다. 따라서 해당 문제를 dfs 방식으로 해결하기 위해서는 stack 제한을 풀어야 한다.

Python
import sys

sys.setrecursionlimit(100000)

굳이 dfs로 해결해야 할 필요는 없기 때문에, bfs 또는 dijkstra로 해결하는 것이 좋은 선택지일 수도 있다.

해당 문제는 bfs, dijkstra 모두 답을 도출해낼 수 있다. bfs의 경우 시간복잡도는 O(V + E) = O(10만 + 50만)이고, dijkstra의 경우 시간복잡도는 O((V + E)logV)이다. 이번에 푸는 방식은 다익스트라 알고리즘을 이용해보기로 한다.

다익스트라 알고리즘을 통해, 강철 부대로부터 각 지점까지의 위치를 구한다.

다익스트라 알고리즘의 자세한 원리는 생략한다.

1. 우선 distance 배열을 INF로 초기화한다. 최댓값으로 초기화 하는 이유는 최솟값을 구해야 하기 때문이다.
2. 우선 시작점을 heap에 넣은 뒤, 노드를 꺼내면서 수행한다. 이때 distance[next]가 distance[node] + 1보다 작다면, 더 이상 탐색을 수행하지 않는다.
3. distnace[next] > distnace[node] + 1이라면, distance[next]를 갱신하고 큐에 삽입한다.

Python
def dijkstra(start, adj, distance):
    q = []
    distance[start] = 0
    heapq.heappush(q, (0, start))

    while q:
        dist, node = heapq.heappop(q)

        # 메모된 거리가 dist보다 짧다면 더 이상 탐색할 필요가 없다.
        if distance[node] < dist:
            continue

        for nxt in adj[node]:
            nxt_distance = dist + 1

            if distance[nxt] > nxt_distance:
                distance[nxt] = nxt_distance
                heapq.heappush(q, (nxt_distance, nxt))

source를 바탕으로 정답을 반환한다.

dijkstra를 성공적으로 실행했다면, distance 배열은 강철 부대로부터 거리가 저장되어 있을 것이다. 이를 source 순서에 맞게 반환하면 된다. 이 때 주의해야 할 점은, 강철 부대로 부터 길이 존재하지 않아 INF의 값을 가진다면 -1을 출력할 수 있도록 바꿔주어야 한다.

통과 코드

배운 것

당연히 다익스트라 알고리즘이 더 빠를 줄 알았는데, 시간 복잡도를 찾아보니까 bfs가 선형적인 시간복잡도를 가지기 때문에 조금 더 빨랐다. 간선의 가중치가 1인 경우에는 bfs를 사용하는 것이 좋은 방식일 수도 있을 듯 하다.