BOJ)13511 트리와 쿼리2

문제: icpc.me/13511

트리에서 2가지 종류의 쿼리를 수행하는 문제이다.

1번 쿼리의 경우 정점들의 거리와 같은 문제로

트리의 루트로 부터 정점 x로의 거리를 dist[x]라고 하였을 때

1번 쿼리의 답은 dist[u]+dist[v]-2*dist[lca(u,v)]가 된다. 왜 이렇게 되는지는 그림을 그려보면 쉽게 알 수 있다.

2번 쿼리는 lca를 구할 때 만든 sparse table을 이용해줄 수 있다.(선형으로 탐색할 경우 시간초과를 피할 수 없다.)

만약 dph[x]-dph[lca]가 z-1보다 크거나 작다면 x를 z-1번 올려주면 되고 (올려줄 때 sparse table을 이용하여 log시간에 올려주어야 한다.)

아닌 경우 두 경로사이에 속한 정점의 수 -z 만큼 y를 올려준다.

#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;
typedef long long ll;
ll n, m, par[100010][22], dist[100010], dph[100010], x, y, z, q;
vector<vector<pair<ll, ll>>> vt;
void dfs(ll here, ll parent, ll d) {
    dph[here] = dph[parent] + 1;
    dist[here] = d;
    for (auto next : vt[here]) {
        if (next.first == parent)continue;
        dfs(next.first, here, d + next.second);
        par[next.first][0] = here;
    }
}
ll getlca(ll x, ll y) {
    if (dph[x] > dph[y])swap(x, y);
    for (int i = 20; i >= 0; i--) {
        if (dph[y] - dph[x] >= (1LL << i))
            y = par[y][i];
    }
    if (x == y)return x;
    for (int i = 20; i >= 0; i--) {
        if (par[x][i] != par[y][i]) {
            x = par[x][i];
            y = par[y][i];
        }
    }
    return par[x][0];
}
int main() {
    scanf("%lld", &n);
    vt.resize(n + 1);
    for (int i = 0; i < n - 1; i++) {
        scanf("%lld%lld%lld", &x, &y, &z);
        vt[x].push_back({ y,z });
        vt[y].push_back({ x,z });
    }
    dfs(1, 0, 0);
    for (int j = 1; j <= 20; j++) {
        for (int i = 1; i <= n; i++) 
            par[i][j] = par[par[i][j - 1]][j - 1];
    }
    scanf("%lld", &m);
    while (m--) {
        scanf("%lld", &q);
        if (q == 1) {
            scanf("%lld%lld", &x, &y);
            ll lca = getlca(x, y);
            printf("%lld\n", dist[x] + dist[y] - 2 * dist[lca]);
        }
        else {
            scanf("%lld%lld%lld", &x, &y, &z);
            ll lca = getlca(x, y);
            ll diffx = dph[x] - dph[lca];
            if (diffx + 1 >= z) {
                z--;
                for (int i = 20; i >= 0; i--) {
                    if ((1LL << i) <= z) {
                        z -= (1 << i);
                        x = par[x][i];
                    }
                }
                printf("%lld\n", x);
            }
            else {
                z = dph[y] - dph[lca] + 1 + diffx - z;
                for (int i = 20; i >= 0; i--) {
                    if ((1LL << i) <= z) {
                        z -= (1 << i);
                        y = par[y][i];
                    }
                }
                printf("%lld\n", y);
            }
        }
    }
    return 0;
}