SPOJ375 Query on a tree 【倍增，在线】

题目链接【http://www.spoj.com/problems/QTREE/】

题意：给出一个包含N(N<=10000)节点的无根树，有多次询问，询问的方式有两种1、DIST  a b 求a->b之间的距离。2、KTH a b k 求a->b链上的第k个节点是谁，。如果输入DONE,结束询问。

思路：首先想到用倍增法可以解决第一种询问，只需要在DFS时候维护一个dis[i]（表示i节点到根节点之间的距离，因为是无根树，我们定义节点1为根），dis(a->b)=dis[a]+dis[b]-dia[lca(a,b)]。第二种询问的解法是首先求出a,b的lca,然后根据lca到a到b的深度判断第k个数在左半枝还是右半枝（可以定义lca为左半枝），然后用倍增的方法求出第k个节点。

#include<vector>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int maxn = 10050;
struct node
{
    int id, len, next;
} E[maxn << 1];
int p[maxn << 1], num;
void init()
{
    memset(p, -1, sizeof(p));
    num = 0;
}
void add(int u, int v, int dis)
{
    E[num].id = u;
    E[num].len = dis;
    E[num].next = p[v];
    p[v] = num++;
}
//------------------------------------------邻接表
int T, N;
int fa[20][maxn];//倍增数组
int dis[maxn], dep[maxn];//距离，深度数组
void DFS(int u, int FA)//DFS求出每个节点的深度，每个节点到根节点的距离，和每个节点跳1步到达的位置
{
    for(int l = p[u]; l != -1; l = E[l].next)
    {
        int id = E[l].id;
        if(id == FA) continue;
        fa[0][id] = u;
        dep[id] = dep[u] + 1;
        dis[id] = dis[u] + E[l].len;
        DFS(id, u);
    }
}
void INIT()//初始化倍增数组
{
    memset(fa, -1, sizeof(fa));
    memset(dis, 0, sizeof(dis));
    memset(dep, 0, sizeof(dep));
    DFS(1, -1);
    for(int k = 0; k + 1 <= 15; k++)
    {
        for(int u = 1; u <= N; u++)
        {
            if(fa[k][u] < 0)
                fa[k + 1][u] = -1;
            else
                fa[k + 1][u] = fa[k][fa[k][u]];
        }
    }
}
int LCA(int u, int v)
{
    if(dep[u] > dep[v]) swap(u, v);
    for(int k = 0; k <= 15; k++)//微调，使得dep相等
        if((dep[v] - dep[u]) >> k & 1)
            v = fa[k][v];
    if(u == v) return u;//公共祖先是u，v中的其中一个
    for(int k = 15; k >= 0; k--)//一起向上跳，先跳距离远的
    {
        if(fa[k][u] != fa[k][v])
        {
            u = fa[k][u];
            v = fa[k][v];
        }
    }
    return fa[0][u];
}
int KTH(int u, int v, int lca, int k)//从u->v路径上的第k个数
{
    if(k == 1)    return u;
    int l1 = dep[u] - dep[lca] + 1;//左半枝
    int l2 = dep[v] - dep[lca];
    if(k <= l1)//在左半枝（lca定义成左半枝）
    {
        k -= 1;
        for(int t = 15; t >= 0; t--)
        {
            if(k >> t & 1)
                u = fa[t][u];
        }
        return u;
    }
    else//在右半枝
    {
        k -= l1;
        k = l2 - k;
        for(int t = 15; t >= 0; t--)
        {
            if(k >> t & 1)
                v = fa[t][v];
        }
        return v;
    }
}
int main ()
{
    scanf("%d", &T);
    while(T--)
    {
        init();
        int u, v, ds, k;
        char s[20];
        scanf("%d", &N);
        for(int i = 1; i < N; i++)
        {
            scanf("%d%d%d", &u, &v, &ds);
            add(u, v, ds);
            add(v, u, ds);
        }
        INIT();
        while(scanf("%s", s))
        {
            if(!strcmp(s, "DONE"))
                break;
            else if(!strcmp(s, "DIST"))
            {
                scanf("%d%d", &u, &v);
                printf("%d
", dis[u] + dis[v] - 2 * dis[LCA(u, v)]);
            }
            else
            {
                scanf("%d%d%d", &u, &v, &k);
                int lca = LCA(u, v);
                printf("%d
", KTH(u, v, lca, k));

            }
        }
    }
    return 0;
}