백준[13511] 트리와 쿼리 2

문제 설명

문제 링크

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

풀이

LCA를 이용하여 풀이할 수 있다. 백준 LCA2를 이해했다고 가정하고 설명하겠다.

1번 출력은 간단하다. dist[i]를 루트로부터 거리라고 할 때 dist[u] + dist[v] - 2*dist[LCA(u, v)]를 구하면 된다. 그림을 그려본다면 쉽게 이해할 수 있을 거다.

2번 출력은 u와 LCA사이에 있는 정점인지 v와 LCA사이에 있는 정점인지 판별한 후 부모 노드로 이동하면 된다.

코드

#include <cstdio>
#include <vector>

using namespace std;
using ll = long long;
using pll = pair<ll, ll>;

vector<pll> v[100005];
ll par[22][100005], lev[100005], dist[100005];

void dfs(ll cur, ll prv) {
    par[0][cur] = prv;
    lev[cur] = lev[prv] + 1;
    for (auto nxt : v[cur]) {
        if (prv != nxt.first) {
            dist[nxt.first] = dist[cur] + nxt.second;
            dfs(nxt.first, cur);
        }
    }
}

ll lca(ll a, ll b) {
    if (lev[a] > lev[b]) swap(a, b);
    for (ll i = 20; i >= 0; i--) {
        if (lev[b] - lev[a] >= (1LL << i)) b = par[i][b];
    }

    if (a == b) return a;

    for (ll i = 20; i >= 0; i--) {
        if (par[i][a] != par[i][b]) {
            a = par[i][a];
            b = par[i][b];
        }
    }
    return par[0][a];
}

int main() {
    ll n; scanf("%lld", &n);
    for (ll i = 0; i < n - 1; i++) {
        ll a, b, c; scanf("%lld %lld %lld", &a, &b, &c);
        v[a].push_back({b, c});
        v[b].push_back({a, c});
    }

    dfs(1, 0);
    for (ll k = 1; k <= 20; k++) {
        for (ll i = 1; i <= n; i++) {
            par[k][i] = par[k - 1][par[k - 1][i]];
        }
    }

    ll m; scanf("%lld", &m);
    while (m--) {
        ll a; scanf("%lld", &a);
        if (a == 1) {
            ll u, v; scanf("%lld %lld", &u, &v);
            printf("%lld\n", dist[u] + dist[v] - 2 * dist[lca(u, v)]);
        } else {
            ll u, v, k; scanf("%lld %lld %lld", &u, &v, &k);
            ll l = lca(u, v);
            if (lev[u] - lev[l] + 1 < k) { 
                k = lev[u] + lev[v] - 2 * lev[l] - k + 1;
                for (ll i = 20; i >= 0; i--) {
                    if (k >= (1LL << i)) {
                        k -= (1LL << i);
                        v = par[i][v];
                    }
                }
                printf("%lld\n", v);
            } else {
                k--;
                for (ll i = 20; i >= 0; i--) {
                    if (k >= (1LL << i)) {
                        k -= (1LL << i);
                        u = par[i][u];
                    }
                }
                printf("%lld\n", u);
            }
        }
    }
}