SPOJ - QTREE2

#include <bits/stdc++.h>
using namespace std;

#define mx 10010
vector<pair<int, int>>adj[mx];
int L[mx], T[mx], P[mx][15], root, cost[mx];
bool vis[mx];

void dfs(int from, int u, int dep, int val)
{
    T[u] = from;
    L[u] = dep;
    cost[u] = val;
    for(auto it : adj[u])
    {
        if(it.first == from)
            continue;
        dfs(u, it.first, dep+1, val+it.second);
    }
}

void lca_init(int n)
{
    memset(P, -1, sizeof(P));
    for(int i = 1; i <= n; i++)
        P[i][0] = T[i];

    for(int j = 1; (1 << j) < n; j++)
        for(int i = 1; i <= n; i++)
        {
            if(P[i][j - 1] != -1)
                P[i][j] = P[P[i][j - 1]][j - 1];
        }
}

int lca_query(int p, int q)
{
    if(L[p] < L[q])
        swap(p, q);
    int log = 1;
    while(1)
    {
        int next = log + 1;
        if((1 << next) > L[p])
            break;
        log++;
    }

    for(int i = log; i >= 0; i--)
        if(L[p] - (1 << i) >= L[q])
            p = P[p][i];
    if(p == q)
        return p;
    for(int i = log; i >= 0; i--)
        if(P[p][i] != -1 && P[p][i] != P[q][i])
            p = P[p][i], q = P[q][i];
    return T[p];
}

int kth(int nd, int k)
{
    int log = log2(L[nd]);
    for(int i = log; i >= 0; i--)
        if(L[nd] - (1 << i) >= k)
            nd = P[nd][i];
    return nd;
}

int main()
{
    int n, p, q, t, i, j, v;
    root = 1;
    scanf("%d", &t);
    while(t--)
    {
        scanf("%d", &n);
        for(i = 1; i < n; i++)
        {
            scanf("%d %d %d", &p, &q, &v);
            adj[p].push_back({q, v});
            adj[q].push_back({p, v});
        }
        dfs(-1, 1, 0, 0);
        lca_init(n);
        char s[10];
        while(1)
        {
            scanf("%s", s);
            if(s[1] == 'O')
                break;
            scanf("%d %d", &p, &q);
            int l = lca_query(p, q);
            if(s[1] == 'I')
            {
                int cp, cq, cl;
                cp = cost[p], cq = cost[q], cl = cost[l];
                printf("%d\n", cp+cq - 2*cl);
            }
            else
            {
                scanf("%d", &v);
                int tt = L[p] - L[l] + 1;
                if(tt >= v)
                    v = kth(p, L[p] - v + 1);
                else
                    v = kth(q, 2*L[l] + v - L[p] - 1);
                printf("%d\n", v);
            }
        }
        puts("");
    }
    return 0;
}