문제

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 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

풀이

아이디어

• 결국 문제가 원하는 건, 각 부대원이 있는 곳 (source의 원소)부터 destination까지의 각각의 최소 거리를 구하라는 말인데, 이는 다르게 생각하면, 말을 source와 destination이라고 해놔서 그렇지 사실상 destination에서 출발해서 각각의 source들까지 도착하는 데 걸리는 최소 비용이라고 해석할 수 있다.
• 그렇게 되면 이 문제는 그냥 다익스트라 그 자체이다. 따라서 출발노드를 destination이라고 해놓고, distance 배열을 다익스트라 돌려서 모두 업데이트 한 뒤, 각 source의 원소에 대해서 distance 배열에 있는 값을 참조해서 answer에 하나씩 append해주면 된다.

코드

from heapq import heappush, heappop
from collections import defaultdict

INF = int(1e9)
def solution(n, roads, sources, destination):
    distance = [INF] * (n+1)
    graph = defaultdict(list)
    
    #그래프 만들기
    for src, dest in roads:
        graph[src].append(dest)
        graph[dest].append(src)
    
    #다익스트라 함수
    def dijkstra(start):
        q = []
        distance[start] = 0
        heappush(q, (0, start))
        
        while q:
            dist, now = heappop(q)
            
            #만약 처리된 노드라면 스킵한다
            if dist > distance[now]:
                continue
            #처리되지 않은 노드라면 다음 과정 수행
            for v in graph[now]:
                cost = dist + 1
                if cost < distance[v]:
                    heappush(q, (cost, v))
                    distance[v] = cost
    
    #다익스트라 수행
    dijkstra(destination)
    
    answer = []
    for s in sources:
        if distance[s] == INF:
            answer.append(-1)
        else:
            answer.append(distance[s])
    
    return answer

C++ 코드

#include <string>
#include <vector>
#include <map>
#include <queue>
#include <algorithm>

using namespace std;
int dest; // 최종목적지
int regionNum; // 지역의 총 개수
const int INF = 1e9;
using pii = pair<int, int>;

int dijkstra(int src, map<int, vector<int>> graph, vector<int>& distance) {
    distance[src] = 0;
    priority_queue<pii, vector<pii>, greater<pii>> heap;
    heap.push(make_pair(0, src)); //(비용, 노드) 순으로 저장
    
    while(!heap.empty()) {
        auto [cost, now] = heap.top();
        heap.pop();
        if(distance[now] < cost) {
            continue;
        }
        //연결된 노드들: graph[now]에 담겨있음
        for(int node: graph[now]) {
            int new_c = 1 + cost;
            if(distance[node] > new_c) {
                distance[node] = new_c;
                heap.push(make_pair(new_c, node));
            }
        }
    }
    
    //최종목적지까지의 최단거리를 리턴
    int ret = (distance[dest] == INF) ? -1 : distance[dest];
    return ret;
}

vector<int> solution(int n, vector<vector<int>> roads, vector<int> sources, int destination) {
    dest = destination;
    regionNum = n;
    //그래프 만들기
    map<int, vector<int>> graph;
    for(int i=0; i<roads.size(); i++) {
        int s = roads[i][0];
        int d = roads[i][1];
        graph[s].push_back(d);
        graph[d].push_back(s);
    }
    
    vector<int> answer;
    vector<int> distance(n+1, INF);
    dijkstra(dest, graph, distance);
    
    for(int node: sources) {
        int ret = distance[node] == INF ? -1 : distance[node];
        answer.push_back(ret);
    }
    
    return answer;
}