【SPOJ】913 Query on a tree II

#include<cstdio>
#include<cstring>
#include<queue>
#include<iostream>
#define MAXN 10010
#define MAXM 20010
using namespace std;
struct LCT {
    int bef[MAXN];
    int next[MAXN][2], pre[MAXN], key[MAXN], sum[MAXN], num[MAXN];
    void Init() {
        memset(next, 0, sizeof(next));
        memset(pre, 0, sizeof(pre));
    }
    inline void PushUp(int x) {
        sum[x] = sum[next[x][0]] + sum[next[x][1]] + key[x];
        num[x] = num[next[x][0]] + num[next[x][1]] + 1;
    }
    inline void Rotate(int x, int kind) {
        int y, z;
        y = pre[x];
        z = pre[y];
        next[y][!kind] = next[x][kind];
        pre[next[x][kind]] = y;
        next[z][next[z][1] == y] = x;
        pre[x] = z;
        next[x][kind] = y;
        pre[y] = x;
        PushUp(y);
    }
    void Splay(int x) {
        int rt;
        for (rt = x; pre[rt]; rt = pre[rt])
            ;
        if (x != rt) {
            bef[x] = bef[rt];
            bef[rt] = 0;
            while (pre[x]) {
                if (next[pre[x]][0] == x)
                    Rotate(x, 1);
                else
                    Rotate(x, 0);
            }
            PushUp(x);
        }
    }
    void Access(int x) {
        int father;
        for (father = 0; x; x = bef[x]) {
            Splay(x);
            pre[next[x][1]] = 0;
            bef[next[x][1]] = x;
            next[x][1] = father;
            pre[father] = x;
            bef[father] = 0;
            father = x;
            PushUp(x);
        }
    }
    int DIST(int x, int y) {
        Access(y);
        for (y = 0; x; x = bef[x]) {
            Splay(x);
            if (!bef[x])
                return sum[y] + sum[next[x][1]];
            pre[next[x][1]] = 0;
            bef[next[x][1]] = x;
            next[x][1] = y;
            pre[y] = x;
            bef[y] = 0;
            y = x;
            PushUp(x);
        }
        return 0;
    }
    int Select(int x, int k) {
        while (num[next[x][0]] + 1 != k) {
            if (num[next[x][0]] + 1 > k)
                x = next[x][0];
            else {
                k -= num[next[x][0]] + 1;
                x = next[x][1];
            }
        }
        return x;
    }
    int KTH(int y, int x, int k) {
        Access(y);
        for (y = 0; x; x = bef[x]) {
            Splay(x);
            if (!bef[x]) {
                if (num[next[x][1]] + 1 == k)
                    return x;
                else if (num[next[x][1]] + 1 > k)
                    return Select(next[x][1], num[next[x][1]] - k + 1);
                return Select(y, k - num[next[x][1]] - 1);
            }
            pre[next[x][1]] = 0;
            bef[next[x][1]] = x;
            next[x][1] = y;
            pre[y] = x;
            bef[y] = 0;
            y = x;
            PushUp(x);
        }
        return 0;
    }
} lct;
int first[MAXN], next[MAXM], v[MAXM], cost[MAXM], e;
bool vis[MAXN];
int INT() {
    int res;
    char ch;
    bool neg;
    while (ch = getchar(), !isdigit(ch) && ch != '-')
        ;
    if (ch == '-') {
        neg = true;
        res = 0;
    } else {
        neg = false;
        res = ch - '0';
    }
    while (ch = getchar(), isdigit(ch))
        res = res * 10 + ch - '0';
    return neg ? -res : res;
}
char CHAR() {
    char ch, res;
    while (ch = getchar(), !isalpha(ch))
        ;
    while (ch = getchar(), isalpha(ch))
        res = ch;
    return res;
}
void AddEdge(int x, int y, int val) {
    v[e] = y;
    cost[e] = val;
    next[e] = first[x];
    first[x] = e++;
}
void BFS(int x) {
    int i, y;
    queue<int> q;
    memset(vis, false, sizeof(vis));
    vis[x] = true;
    q.push(x);
    while (!q.empty()) {
        x = q.front();
        q.pop();
        for (i = first[x]; i != -1; i = next[i]) {
            y = v[i];
            if (!vis[y]) {
                lct.bef[y] = x;
                lct.key[y] = cost[i];
                vis[y] = true;
                q.push(y);
            }
        }
    }
}
int main() {
    char ch;
    int c, n, i, x, y, k;
    c = INT();
    while (c--) {
        n = INT();
        lct.Init();
        memset(first, -1, sizeof(first));
        for (e = 0, i = 1; i < n; i++) {
            x = INT(), y = INT(), k = INT();
            AddEdge(x, y, k);
            AddEdge(y, x, k);
        }
        BFS(1);
        while (ch = CHAR(), ch != 'E') {
            x = INT(), y = INT();
            if (ch == 'T')
                printf("%d\n", lct.DIST(x, y));
            else {
                k = INT();
                printf("%d\n", lct.KTH(x, y, k));
            }
        }
        putchar('\n');
    }
    return 0;
}