【LCA】SPOJ QTREE2

#include <cstdio>
#include <iostream>
#include <cstdlib>
#include <memory.h>
#include <ctime>
#include <algorithm>
#include <cmath>
#include <cstring>
#include <string>
#include <utility>
#include <stack>
#include <set>
#include <map>
#include <queue>
#include <vector>
#include <climits>
using namespace std;

#define clr(_arr) (memset(_arr, 0, sizeof(_arr)))

typedef long long ll;

const int INF = INT_MAX,
          N_MAX  = 10010,
          LG_N_MAX = 20;

int N, LG_N,
    a, b, c,
    root = 1,
    par[LG_N_MAX][N_MAX],//par[m][u] indicates the (2 ^ m)th ancestor of u
    dep[N_MAX], disRt[N_MAX];

char instr[8];

vector<int> G[N_MAX];

typedef pair<int, int> pii;

map<pii, int> len;

void init()
{
    clr(par);
    clr(dep);
    clr(disRt);

    for (int i = 1; i <= N; ++i)
        G[i].clear();

    len.clear();

    LG_N = 0;
    int _N = 1;

    while (_N <= N)
    {
        _N <<= 1;
        ++LG_N;
    }
    ++LG_N;
}

void addEdge(int a, int b, int c)
{
    G[a].push_back(b);
    G[b].push_back(a);

    len[make_pair(a, b)] = c;
    len[make_pair(b, a)] = c;
}

void dfs(int u, int p, int d, int dR)
{
    par[0][u] = p;
    dep[u] = d;
    disRt[u] = dR;

    int out = G[u].size();
    for (int i = 0; i < out; ++i)
    {
        int s = G[u][i];

        if (s != p)
            dfs(s, u, d + 1, dR + len[pii(u, s)]);
    }
}

void pre()
{
    dfs(root, -1, 0, 0);

    //par[m + 1][u] can be obtained only after all par[m][v](v = 1..N) is obtained, and thus the outer loop should be of u
    for (int u = 1; u <= N; ++u)
        for (int m = 0; m + 1 <= LG_N; ++m)
        {
            if (par[m][u] < 0)
                par[m + 1][u] = -1;
            else
                par[m + 1][u] = par[m][par[m][u]];
        }
}

int kthAnt(int u, int k)
{
    for (int m = 0; k; ++m, k >>= 1)
    {
        if (k & 1)
        u = par[m][u];
    }

    return u;
}

int LCA(int u, int v)
{
    if (dep[u] < dep[v])
        swap(u, v);

    u = kthAnt(u, dep[u] - dep[v]);

    if (u == v)
        return u;

    //binary search, search all ancestors of u and v for the last different node index pair, whose parent is LCA(u, v)
    for (int m = LG_N; m >= 0; --m)
        if (par[m][u] != par[m][v])
        {
            u = par[m][u];
            v = par[m][v];
        }

    return par[0][u];
}

int Kth(int u, int v, int k)
{
    int w = LCA(u, v);

    int uw = dep[u] - dep[w] + 1,
        wv = dep[v] - dep[w] + 1,
        _k = k;

    k = (_k <= uw) ? k : (wv - (k - uw));
    u = (_k <= uw) ? u : v;

    return kthAnt(u, k - 1);
}

int main()
{
    int t;

    scanf("%d", &t);

    while(t--)
    {
        bool done = false;

        scanf("%d", &N);

        init();

        for (int i = 1; i < N; ++i)
        {
            scanf("%d %d %d", &a, &b, &c);

            addEdge(a, b, c);
        }

        pre();

        getchar();
        while(!done)
        {
            scanf("%s", instr);

            switch (instr[1])
            {
                case ‘I‘:
                {
                    scanf("%d %d\n", &a, &b);

                    int w = LCA(a, b);

                    int dis = disRt[a] + disRt[b] - 2 * disRt[w];

                    printf("%d\n", dis);

                    break;
                }
                case ‘T‘:
                {
                    int k;
                    scanf("%d %d %d\n", &a, &b, &k);

                    printf("%d\n", Kth(a, b, k));

                    break;
                }
                case ‘O‘:
                {
                    done = true;

                    break;
                }
            }
        }
        printf("\n");
    }

    return 0;
}