백준[1761] 정점들의 거리

문제 설명

문제 링크

http://icpc.me/1761

1761번: 정점들의 거리

풀이

LCA(최소 공통 조상)를 구하여 풀 수 있다.

root에서 a까지의 거리를 cost[a]라고 할 때, 트리에서 정점 a, b의 거리는 cost[a] + cost[b] - 2 * cost[lca]라고 할 수 있다.

코드

#include <iostream>
#include <algorithm>
#include <vector>

using namespace std;

const int N = 4e4 + 4;
vector<pair<int, int>> v[N];
int dp[22][N], cost[N], lev[N];
bool visited[N];

void dfs(int cur, int prv) {
    visited[cur] = true;
    dp[0][cur] = prv;
    lev[cur] = lev[prv] + 1;
    for (auto [nxt, c] : v[cur]) {
        if (!visited[nxt]) {
            cost[nxt] = cost[cur] + c;
            dfs(nxt, cur);
        }
    }
}

int lca(int a, int b) {
    if (lev[a] > lev[b]) swap(a, b);
    for (int i = 20; i >= 0; i--) {
        if (lev[b] - lev[a] >= (1 << i)) b = dp[i][b];
    }
    if (a == b) return a;

    for (int i = 20; i >= 0; i--) { 
        if (dp[i][a] != dp[i][b]) {
            a = dp[i][a];
            b = dp[i][b];
        }
    }
    return dp[0][a];
    // return 1;
}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n; cin >> n;
    for (int i = 0; i < n - 1; i++) {
        int a, b, c; cin >> a >> b >> c;
        v[a].push_back({b, c});
        v[b].push_back({a, c});
    }
    dfs(1, 0);

    for (int i = 1; i <= 20; i++)
        for (int j = 1; j < N; j++)
            dp[i][j] = dp[i - 1][dp[i - 1][j]];
    
    int q; cin >> q;
    while (q--) {
        int a, b; cin >> a >> b;
        int _lca = lca(a, b);
        // cout << _lca << endl;
        int ret = cost[a] + cost[b] - 2 * cost[_lca];   
        cout << ret << '\n';
    }
}