SPOJ - QTREE2 Query on a tree II 树上倍增LCA

传送门
题意：给定树和边权，在树上实现两个操作
1、求x，y两个节点之间的距离(DIST 4 6)
2、求从x出发，y为终点路径上的第k个节点(KTH 4 6 1，就是4)
操作1:dist[x] + dist[y] - 2 * dist[lca]

操作2:先求x到y的链长dep[x] + dep[y] - 2 * dep[lca]，然后分情况看第k个点在x-lca链还是lca-y链，然后求祖先节点即可。

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int N = 2e5 + 100;
int n, x, y, cost, tot;
int Next[N], ver[N], edge[N], dist[N], dep[N], fa[N][50], head[N];

void init(){
    tot = 0;
    memset(fa, 0, sizeof(fa));
    memset(dep, 0, sizeof(dep));
    memset(Next, 0, sizeof(Next));
    memset(dist, 0, sizeof(dist));
    memset(head, 0, sizeof(head));
}

void add(int x, int y, int z){
    ver[++tot] = y;edge[tot] = z;Next[tot] = head[x];head[x] = tot;
}


void dfs(int x, int f){
    for(int i = 1; i < 30; i++)
        if(fa[fa[x][i - 1]][i - 1]) fa[x][i] = fa[fa[x][i - 1]][i - 1];
        else break;
    for(int i = head[x]; i; i = Next[i]){
        int to = ver[i];
        if(to == f) continue;
        fa[to][0] = x;
        dist[to] = dist[x] + edge[i];
        dep[to] = dep[x] + 1;
        dfs(to, x);
    }
}


int getlca(int x, int y){
    if(dep[x] < dep[y]) swap(x, y);
    int cnt = dep[x] - dep[y];
    for(int i = 0; i < 30; i++)
        if((1 << i) & cnt) x = fa[x][i];
    if(x == y) return x;
    for(int i = 30; i >= 0; i--){
        if(fa[x][i] != fa[y][i]){
            x = fa[x][i];
            y = fa[y][i];
        }
    }
    return fa[x][0];
}


int find(int x, int k){
    for(int i = 0; i < 30; i++){
        if((1<<i) & k) x = fa[x][i];
    }
    return x;
}

int main(){
    ios::sync_with_stdio(false);
    int t;
    cin >> t;
    while(t--){
        init();
        cin >> n;
        for(int i = 1; i < n; i++){
            cin >> x >> y >> cost;
            add(x, y, cost);
            add(y, x, cost);
        }
        dfs(1, 0);
        string s;
        while(cin >> s){
            if(s == "DIST"){
                int a, b, lca;
                cin >> a >> b;
                lca = getlca(a, b);
                int ans = dist[a] + dist[b] - dist[lca] * 2;
                cout << ans << endl;
            }
            else if(s == "KTH"){
                int a, b, lca, kth;
                cin >> a >> b >> kth;
                lca = getlca(a, b);
                if(kth <= dep[a] - dep[lca] + 1) cout << find(a, kth - 1) << endl;
                else{
                    int tmp = dep[a] + dep[b] - dep[lca] * 2 + 1;
                    cout << find(b, tmp - kth) << endl;
                }
            }
            else break;
        }
        
        cout << endl;
    }
    return 0;
}