백준 1761번: 정점들의 거리 (C++)

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

1761번: 정점들의 거리

---

Solved By: LCA

 

Sparse Table을 이용해 LCA만 구할 수 있었지만, 각 정점들 간의 거리 또한 Sparse Table로 parent node들을 구하듯이 거리를 구할 수 있습니다.

 

parent를 구하는 점화식과 유사합니다. Sparse Table을 아는 상태에서 조금 더 생각해보면 충분히 나올 수 있는 점화식입니다.

 

#include <iostream>
#include <stack>
#include <vector>
#include <algorithm>
#define MAX 40001
#define LOG_MAX 16
#define pii pair<int, int>
using namespace std;

int n, k;
int parent[MAX][LOG_MAX];
int dist[MAX][LOG_MAX];
int level[MAX];
vector<pii> adj[MAX];

void dfs(int now, int par){
    for(int i = 0; i < adj[now].size(); i++){
        int next = adj[now][i].first;
        int cost = adj[now][i].second;
        
        if(next == par) continue;
        parent[next][0] = now;
        level[next] = level[now] + 1;
        dist[next][0] = cost;
        dfs(next, now);
    }
}

void swap(int *a, int *b){
    int* temp = a;
    a = b;
    b = temp;
}

int lca(int a, int b){
    int res = 0;
    int maxV = 0, minV = 1e9;
    if(level[a] < level[b]) swap(a, b);
    int diff = level[a] - level[b];
    
    for(int i = LOG_MAX - 1; i >= 0; i--){
        if(diff >= 1 << i){
            diff -= 1 << i;
            res += dist[a][i];
            a = parent[a][i];
        }
    }
    
    if(a != b){
        for(int i = LOG_MAX - 1; i >= 0; i--){
            if(parent[a][i] != 0 && parent[a][i] != parent[b][i]){
                res += (dist[a][i] + dist[b][i]);
                
                a = parent[a][i];
                b = parent[b][i];
            }
        }
        
        res += dist[a][0];
        res += dist[b][0];
    }
    
    
    return res;
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    cin >> n;
    for(int i = 0; i < n - 1; i++){
        int a, b, c; cin >> a >> b >> c;
        adj[a].push_back({b, c});
        adj[b].push_back({a, c});
    }
    
    dfs(1, 0);
    
    for(int j = 1; j < LOG_MAX; j++){
        for(int i = 1; i <= n; i++){
            parent[i][j] = parent[parent[i][j - 1]][j - 1];
            dist[i][j] = dist[parent[i][j - 1]][j - 1] + dist[i][j - 1];
        }
    }
    
    cin >> k;
    while(k--){
        int a, b; cin >> a >> b;
        cout << lca(a, b) << '\n';
    }
    
    return 0;
}