[프로그래머스]level.3 - 부대복귀 (C++)

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

문제 설명

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

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

내 풀이

다익스트라 알고리즘을 사용해 목적지(destination)부터 모든 지점까지의 거리를 구한 뒤 조건에 맞게 반환해주면 된다.

#include <string>
#include <vector>
#include <queue>

using namespace std;

vector<int> solution(int n, vector<vector<int>> roads, vector<int> sources, int destination) {
    vector<int> answer;
    vector<vector<int>> map(n);
    int max_int = 999'999;
    vector<int> distance(n, max_int);
    priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int,int>>> pq;
    
    for(const auto& r : roads)
    {
        map[r[0]-1].push_back(r[1]-1);
        map[r[1]-1].push_back(r[0]-1);
    }
    
    // 목적지를 큐에 추가
    pq.push({0, destination-1});
    
    while(!pq.empty())
    {
        auto top = pq.top();
        pq.pop();
        
        int r = top.second;
        int d = top.first;
        
        distance[r] = min(distance[r], d);
        
        for(int i = 0; i < map[r].size(); ++i)
        {   
            int nxt = map[r][i];
            int dist = d + 1;
            
            if(dist < distance[nxt])
            {
                pq.push({dist, nxt});
            }
        }
    }
    
    for(int s : sources)
    {
        if(max_int != distance[s-1])
            answer.push_back(distance[s-1]);
        else
            answer.push_back(-1);
    }
    
    return answer;
}

채점 결과