백준 1240 노드사이의 거리 C++

 

https://www.acmicpc.net/problem/1240

1240번: 노드사이의 거리

처음에 문제를 보고 엣지 길이가 항상 1일거라 생각하고 예제를 봤는데 도저히 나올 수 없는 트리 구조여서 어떻게 하나 막막했던 문제였습니다. (하지만 문제를 다시 읽어보니 엣지 가중치가 있었다.....)

 

1000개 노드를 대상으로 탐색을 진행하고, 1000개 쿼리가 발생할 수 있으므로 그냥 다익스트라로 접근해도 뭔가 풀릴 거 같지만, 안전하게 우선순위 큐를 이용한 다익스트라로 문제를 풀었습니다. 

 

방식은 간단합니다. 입력으로 주어진 노드를 대상으로 그래프를 그리고 다익스트라 알고리즘 + 우선순위 큐 이용한 방식으로 해결하면 노드 사이의 거리를 구할 수 있습니다. 

 

다익스트라 알고리즘을 진행하면서 최단거리 갱신을 우선순위 큐를 이용해 진행하므로 O(N log N) (트리니까 엣지 전체 개수 = 노드 개수 - 1이고 노드 개수 N개 잡았을 때)에 해결할 수 있습니다. 그리고 쿼리도 M개 진행되므로 총 O(NM log N)에 해결할 수 있습니다.

 

코드는 다음과 같습니다.

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>
#include <cstring>

using namespace std;

int n, m;
vector<pair<int, int>> vec[1001]; // {연결된 노드, 가중치}
int dist[1001]; // 거리 최대 1000만

void init()
{
    cin >> n >> m;
    int f, t, w; // from, to, weight
    for(int i = 0; i < (n-1); i++) {
        cin >> f >> t >> w;
        vec[f].push_back({t, w});
        vec[t].push_back({f, w});
    }
}

void dist_clear()
{
    for(int i = 0; i < 1001; i++) {
        dist[i] = 987654321;
    }
}

void dijkstra(int start, int last)
{
    // 우선순위 큐 이용한 다익스트라
    dist[start] = 0;
    
    // {현재까지 거리, 현재 탐색 노드}
    priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>> > pq; 
    pq.push({0, start});
    
    while(!pq.empty()) {
        int d = pq.top().first;
        int cur = pq.top().second;
        pq.pop();
        
        // 계산했던 값이 더 작은 경우 (더 안 해도 됨)
        if(dist[cur] < d)
            continue;
            
        // 탐색 
        for(auto v : vec[cur]){
            int next = v.first;
            int nd = v.second + d;
            
            // 갱신
            if(dist[next] > nd) {
                dist[next] = nd;
                pq.push({nd, next});
            }
        }
    }
    
    cout << dist[last] << '\n';
}

int main()
{
    ios_base :: sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    
    init();
    
    int start, last;
    for(int i = 0; i < m; i++) {
        dist_clear();
        cin >> start >> last;
        dijkstra(start, last);
    }
}