부대복귀 [Lv. 3] (C++)

프로그래머스

강철부대 대원들은 여러 지역으로 작전을 떠났습니다.

이 때 모든 대원들을 특정 지역으로 호출할 때 각 대원별 이동 시간을 구해주세요

적군에 의해 사라진 길이 있어 돌아오지 못하는 대원들도 있습니다.

 

풀이 방법

이 문제는 하나의 트릭이 있습니다.

마치 모든 대원에 대해서 길찾기 알고리즘을 수행해야 할 것 처럼

여러 출발점이 있다고 말합니다. (플로이드-워셜을 생각하게끔 유도합니다)

 

하지만 목적지는 하나이고, 모든 길이 양방향 통행이 가능하기에

굳이 모든 대원에 대해서 탐색할 필요가 없고

목적지에서부터 다익스트라 알고리즘을 한번만 수행하면 

목적지 -> 모든 대원의 위치까지의 거리를 알 수 있게됩니다.

 

이후에 그 거리만 배열에 담아서 출력하면 쉽게 해결할 수 있습니다.

 

정답 코드

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

using namespace std;

void dijkstra(vector<vector<int>>& roads, vector<int>& dists, int start)
{
    dists[start] = 0;
    
    priority_queue<pair<int, int>> q;
    q.push({0, start});
    
    while(!q.empty())
    {
        auto [cost, node] = q.top(); q.pop();
        cost *= -1;
        
        for (int& next : roads[node])
        {
            if (dists[next] == -1 || dists[next] > cost + 1)
            {
                dists[next] = cost + 1;
                q.push({-(cost+1), next});
            }
        } 
    }
}

vector<int> solution(int n, vector<vector<int>> roads, vector<int> sources, int destination) {
    vector<vector<int>> _roads(n+1);
    for (auto& r : roads)
    {
        _roads[r[0]].push_back(r[1]);
        _roads[r[1]].push_back(r[0]);
    }
    
    vector<int> dists(n + 1, -1);
    dijkstra(_roads, dists, destination);
        
    for(int& s : sources)
        s = dists[s];
    
    return sources;
}