SPOJ 913 Query on a tree II（动态树）

题目链接：http://www.spoj.com/problems/QTREE2/

题意：给出一棵树，每条边有一个权值。两种操作：（1）询问某两个节点之间的距离；（2）询问x到y的路径上的第K个节点的编号。

思路：节点保存dis值、cnt。对于第二种询问，先access(x)，然后y向上，父节点为空时，判断第K个节点在y的左侧右侧还是y本身。

struct node

{

    int cnt,dis,val,id;

    node *c[2],*f;

};



node a[N],*nullNode;

int n,m;



void pushUp(node *p)

{

    p->dis=p->c[0]->dis+p->c[1]->dis+p->val;

    p->cnt=p->c[0]->cnt+p->c[1]->cnt+1;

}







void rotate(node *p,int k)

{

    node *q=p->f;

    q->c[k]=p->c[!k];

    if(p->c[!k]!=nullNode) p->c[!k]->f=q;

    p->c[!k]=q;

    p->f=q->f;

    if(q->f!=nullNode)

    {

        if(p->f->c[0]==q) p->f->c[0]=p;

        if(p->f->c[1]==q) p->f->c[1]=p;

    }

    q->f=p;

    pushUp(q);

}



int isRoot(node *p)

{

    return p->f==nullNode||p->f->c[0]!=p&&p->f->c[1]!=p;

}





void splay(node *p)

{

     while (!isRoot(p))

     {

         if(isRoot(p->f))

         {

             if(p==p->f->c[0]) rotate(p,0);

             else rotate(p,1);

         }

         else

         {

             if(p->f->f->c[0]==p->f)

             {

                 if(p->f->c[0]==p) rotate(p->f,0);

                 else rotate(p,1);

                 rotate(p,0);

             }

             else

             {

                 if(p->f->c[1]==p) rotate(p->f,1);

                 else rotate(p,0);

                 rotate(p,1);

             }

         }

     }

     pushUp(p);

}



node* access(node *p)

{

    node *q=nullNode;

    while(p!=nullNode)

    {

        splay(p);

        p->c[1]=q;

        pushUp(p);

        q=p;

        p=p->f;

    }

    return q;

}







int getDis(int x,int y)

{

    access(a+x);

    node *lca=access(a+y);



    splay(a+x);

    if(a+x!=lca) return lca->c[1]->dis+a[x].dis;

    return lca->c[1]->dis;

}





int selectK(node *p,int k)

{

    while(p->c[0]->cnt+1!=k)

    {

        if(p->c[0]->cnt+1>k) p=p->c[0];

        else

        {

            k-=p->c[0]->cnt+1;

            p=p->c[1];

        }

    }

    return p->id;

}



int getKth(int x,int y,int k)

{

    access(a+x);

    node *lca=access(a+y);

    splay(a+x);

    if(a+x==lca)

    {

        if(k==1) return lca->id;

        return selectK(lca->c[1],k-1);

    }

    else

    {

        if(a[x].cnt+1==k) return lca->id;

        if(a[x].cnt+1>k) return selectK(a+x,a[x].cnt+1-k);

        return selectK(lca->c[1],k-1-a[x].cnt);

    }

}



struct Node

{

    int v,dis,next;

};



Node edges[N];

int head[N],e;





void Add(int u,int v,int dis)

{

    edges[e].v=v;

    edges[e].dis=dis;

    edges[e].next=head[u];

    head[u]=e++;



    edges[e].v=u;

    edges[e].dis=dis;

    edges[e].next=head[v];

    head[v]=e++;

}



void build()

{

    queue<int> Q;

    Q.push(1);

    int i,u,v;

    while(!Q.empty())

    {

        u=Q.front();

        Q.pop();



        for(i=head[u];i!=-1;i=edges[i].next)

        {

            v=edges[i].v;

            if(a+v==a[u].f) continue;

            a[v].f=a+u;

            a[v].val=edges[i].dis;

            Q.push(v);

        }

    }



}



int get()

{

    int x=0,flag=1;

    char c=getchar();

    while(!isdigit(c)) c=getchar();

    if(c=='-')

    {

        flag=-1,c=getchar();

        while(!isdigit(c)) c=getchar();

    }

    while(isdigit(c))

    {

        x=x*10+c-'0';

        c=getchar();

    }

    return x*flag;

}







void init()

{

    nullNode=new node();

    nullNode->cnt=0;

    nullNode->dis=0;

    nullNode->c[0]=nullNode->c[1]=nullNode->f=nullNode;

    int i,u,v,dis;

    clr(head,-1);

    e=0;

    FOR1(i,n-1) u=get(),v=get(),dis=get(),Add(u,v,dis);

    FOR1(i,n)

    {

        a[i].cnt=1;

        a[i].dis=0;

        a[i].val=0;

        a[i].id=i;

        a[i].c[0]=a[i].c[1]=a[i].f=nullNode;

    }

    build();

}



int main()

{

    int C=get();

    while(C--)

    {

        n=get();

        init();

        int u,v,K;

        char op[10];

        while(1)

        {

            scanf("%s",op);

            if(op[1]=='O') break;

            else if(op[0]=='D') u=get(),v=get(),PR(getDis(u,v));

            else u=get(),v=get(),K=get(),PR(getKth(u,v,K));

        }

        puts("");

    }

    return 0;

}