SPOJ375 QTREE - Query on a tree

#include<iostream>
#include<cstdio>
#include<cstring>
#define lc o<<1
#define rc o<<1|1
using namespace std;
const int N=10000+5;
const int inf=1e9;
struct Node{
int l,r,mx;
}tr[N*4];
struct Edge{int to,next,v;}e[N*2];
int head
,cnt,d
[2];
int dep
,siz
,son
,fa
,top
,pos
,sz,cost
,rank
;
void ins(int u,int v,int w){
e[++cnt].to=v;e[cnt].next=head[u];head[u]=cnt;e[cnt].v=w;
}
void dfs1(int u){
siz[u]=1;son[u]=0;
for(int i=head[u];i;i=e[i].next){
int v=e[i].to;
if(v==fa[u])continue;
fa[v]=u;dep[v]=dep[u]+1;
cost[v]=e[i].v;
dfs1(v);
siz[u]+=siz[v];
if(siz[v]>siz[son[u]])son[u]=v;
}
}
void dfs2(int u,int tp){
pos[u]=++sz;top[u]=tp;rank[sz]=u;
if(son[u])dfs2(son[u],top[u]);
for(int i=head[u];i;i=e[i].next){
int v=e[i].to;
if(v!=fa[u]&&v!=son[u])
dfs2(v,v);
}
}
inline void pushup(int o){tr[o].mx=max(tr[lc].mx,tr[rc].mx);}
void update(int o,int p,int v){
int l=tr[o].l,r=tr[o].r;
if(l==r)tr[o].mx=v;
else{
int m=l+r>>1;
if(p<=m)update(lc,p,v);
else update(rc,p,v);
pushup(o);
}
}
int query(int o,int a,int b){
int l=tr[o].l,r=tr[o].r;
if(l==a&&b==r)return tr[o].mx;
else{
int m=l+r>>1;
if(b<=m)return query(lc,a,b);
else if(m<a)return query(rc,a,b);
else return max(query(lc,a,m),query(rc,m+1,b));
}
}
void build(int o,int l,int r){
tr[o].l=l;tr[o].r=r;
if(l==r)tr[o].mx=cost[rank[l]];
else{
int m=l+r>>1;
build(lc,l,m);build(rc,m+1,r);
pushup(o);
}
}
int ask(int u,int v){
int ans=-inf;
while(top[u]!=top[v]){
if(dep[top[u]]<dep[top[v]])swap(u,v);
ans=max(ans,query(1,pos[top[u]],pos[u]));
u=fa[top[u]];
}
if(u==v)return ans;
if(dep[u]>dep[v])swap(u,v);
return max(ans,query(1,pos[son[u]],pos[v]));
}
int main(){
int T;scanf("%d",&T);
while(T--){
int n;scanf("%d",&n);
memset(head,0,sizeof(head));cnt=0;sz=0;
int u,v,w;
for(int i=1;i<n;i++){
scanf("%d%d%d",&u,&v,&w);
d[i][0]=u;d[i][1]=v;
ins(u,v,w);ins(v,u,w);
}
dfs1(1);dfs2(1,1);build(1,1,n);
for(int i=1;i<n;i++)
if(dep[d[i][0]]<dep[d[i][1]])swap(d[i][0],d[i][1]);
char opt[10];
while(true){
scanf("%s",opt);
if(opt[0]=='D')break;
if(opt[0]=='Q'){
scanf("%d%d",&u,&v);
printf("%d\n",ask(u,v));
}else{
scanf("%d%d",&u,&w);
update(1,pos[d[u][0]],w);
}
}
}
return 0;
}