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;
}