SPOJ-QTREE2 Query on a tree II(暴力＋LCA)

题目大意：给出一棵树，3种操作
DIST u,v 询问u到v的距离
KTH k, u, v 询问u到v的路径上的第k大的边的权值

解题思路：刚开始以为会爆，结果发现不会
直接暴力存储u到v的路上的所有边，再进行排序，输出第k大的边即可

#include <cstdio>
#include <cstring>

#define N 10010

struct Edge{
    int to, next, cost;
}E[2*N];

int head[N], depth[2 * N], first[N], ver[2 * N], pre[N], dis[N], dp[2*N][63];
int n, tot;
bool vis[N];
char str[100];

void AddEdge(int u, int v, int c) {
    E[tot].to = v; E[tot].next = head[u]; E[tot].cost = c; head[u] = tot++;
    u = u ^ v; v = u ^ v; u = u ^ v;
    E[tot].to = v; E[tot].next = head[u]; E[tot].cost = c; head[u] = tot++;
}

void init() {
    scanf("%d", &n);
    memset(head, -1, sizeof(head));
    tot = 0;

    int u, v, c;
    for (int i = 1; i < n; i++) {
        scanf("%d%d%d", &u, &v, &c);
        AddEdge(u, v, c);
    }
}

void dfs(int u, int dep) {
    vis[u] = true; ver[++tot] = u; depth[tot] = dep; first[u] = tot;
    for (int i = head[u]; ~i; i = E[i].next) {
        int v = E[i].to;
        if (!vis[v]) {
            pre[v] = u;
            dis[v] = dis[u] + E[i].cost;
            dfs(v, dep + 1);
            ver[++tot] = u; depth[tot] = dep;
        }
    }
}

void RMQ() {
    for (int i = 1; i <= tot; i++)
        dp[i][0] = i;

    int a, b;
    for (int j = 1; (1 << j) <= tot; j++)
        for (int i = 1; i + (1 << j) - 1 <= tot; i++) {
            a = dp[i][j - 1];
            b = dp[i + (1 << (j - 1))][j-1];
            if (depth[a] < depth[b])
                dp[i][j] = a;
            else
                dp[i][j] = b;
        }
}

int Query(int x, int y) {
    int k = 0;
    while (1 << (k + 1) <= (y - x + 1)) k++;
    int a = dp[x][k];
    int b = dp[y - (1 << k) + 1][k];
    if (depth[a] < depth[b])
        return a;
    return b;
}

int LCA(int u, int v) {
    int a = first[u];
    int b = first[v];
    if (a > b) {
        a = a ^ b; b = a ^ b; a = a ^ b;
    }
    int c = Query(a, b);
    return ver[c];
}

int num[N];
int KTH(int u, int v, int c) {
    int lca = LCA(u, v);
    int cnt = 1;

    int t = u;
    while (t != lca) {
        t = pre[t];
        cnt++;
        if (cnt == c)
            return t;
    }
    c = c - cnt;
    t = v;
    int cnt2 = 0;

    while (t != lca) {
        num[cnt2] = t;
        cnt2++;
        t = pre[t];
    }
//  printf("cnt2 is %d c is %d \n", cnt2, c);
    return num[cnt2 - c];
}

void solve() {
    memset(vis, 0, sizeof(vis));
    tot = 0;
    dis[1] = 0;
    pre[1] = 1;
    dfs(1,1);
    RMQ();
    int u, v, c;
    while (scanf("%s", str)) {
        if (str[0] == 'D' && str[1] == 'O')
            break;
        if (str[0] == 'D') {
            scanf("%d%d", &u, &v);
            int lca = LCA(u, v);
            printf("%d\n", dis[u] + dis[v] - 2 * dis[lca]);
        }
        else {
            scanf("%d%d%d", &u, &v, &c);
            if (c == 1) 
                printf("%d\n", u);
            else 
                printf("%d\n", KTH(u, v, c));
        }
    }

}

int main() {
    int test;
    scanf("%d", &test);
    while (test--) {
        init();
        solve();
    }
    return 0;
}