프로그래머스 - 등산코스 정하기 [C++]

이 문제는 다익스트라 알고리즘을 이용해서 풀었다.

 

처음에는 출발지들에 대해 루프를 돌면서 산봉우리들에 대해 루프를 돌고, 해당 출발지, 산봉우리 쌍에 대해 다익스트라 알고리즘을 사용했다. 그러나 이 방법으로는 당연히 시간 초과가 난다.

 

다른 방법으로 처음 다익스트라 알고리즘을 시작할 때, 우선순위 큐에 모든 출발지를 넣고 각 노드가 SUMMIT을 만나면 종료하도록 구현하는 것이다. 이 방법을 사용하면 한 번의 다익스트라 호출로 산봉우리들로 가는 최단 거리를 알 수 있다.

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

using namespace std;

struct compare
{
    bool operator() (pair<int, int>& left, pair<int, int>& right)
    {
        return left.second > right.second;
    }
};

#define MAX 50000

#define SHELTER 0
#define GATE 1
#define SUMMIT 2
#define INF 100000000

vector<pair<int, int>> edges[MAX + 1];
int info[MAX + 1];

vector<int> starts;
vector<int> tops;
int N;

pair<int, int> dijkstra()
{
    priority_queue<pair<int, int>, vector<pair<int, int>>, compare> pq;
    vector<int> dist(MAX + 1, INF);
    
    for (int start : starts)
    {
        dist[start] = 0;
        pq.push({start, 0});
    }
    
    while (!pq.empty())
    {
        int present = pq.top().first;
        int cost = pq.top().second;
        pq.pop();
        
        if (info[present] == SUMMIT)
        {
            continue;
        }
        
        for (pair<int, int>& edge : edges[present])
        {
            int next = edge.first;
            int nextCost = max(cost, edge.second);
            
            if (dist[next] > nextCost)
            {
                if (info[next] == GATE)
                {
                    continue;
                }
                
                dist[next] = nextCost;
                    
                pq.push({next, nextCost});
            }
        }
    }
    
    int result = INF;
    int topNum = -1;
    
    for (int top : tops)
    {
        if (dist[top] < result)
        {
            result = dist[top];
            topNum = top;
        }
    }
    
    return {topNum, result};
}

vector<int> solution(int n, vector<vector<int>> paths, vector<int> gates, vector<int> summits) {
    vector<int> answer;
    N = n;
    
    sort(summits.begin(), summits.end());
    
    for (vector<int> path : paths)
    {
        int a = path[0];
        int b = path[1];
        int w = path[2];
        
        edges[a].push_back({b, w});
        edges[b].push_back({a, w});
    }
    
    for (int gate : gates)
    {
        info[gate] = GATE;
        starts.push_back(gate);
    }
    
    for (int summit : summits)
    {
        info[summit] = SUMMIT;
        tops.push_back(summit);
    }
    
    pair<int, int> result = dijkstra();
    
    answer.push_back(result.first);
    answer.push_back(result.second);
    
    return answer;
}