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
,*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
;
int head
,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;
}