SPOJ QTREE Query on a tree --树链剖分

#include <iostream>
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <cmath>
#include <algorithm>
using namespace std;
#define N 10007

int siz[N];  //子树大小
int son[N];  //重儿子
int dep[N];  //深度
int pos[N];  //在线段树中的位置
int Top[N];  //所在重链的祖先
int fa[N];   //父节点
int ans[N];  //答案
int head[2*N],tot,POS,n,m;
struct Edge
{
    int v,next;
}G[2*N];
int tree[4*N];
struct node
{
    int u,v,w;
}edge[N];

void init()
{
    POS = tot = 0;
    memset(head,-1,sizeof(head));
    memset(son,-1,sizeof(son));
    memset(tree,0,sizeof(tree));
}

void addedge(int u,int v)
{
    G[tot].v = v, G[tot].next = head[u], head[u] = tot++;
    G[tot].v = u, G[tot].next = head[v], head[v] = tot++;
}

void pushup(int rt)
{
    tree[rt] = max(tree[2*rt],tree[2*rt+1]);
}

void update(int l,int r,int pos,int val,int rt)
{
    if(l == r)
    {
        tree[rt] = val;
        return;
    }
    int mid = (l+r)/2;
    if(pos <= mid)
        update(l,mid,pos,val,2*rt);
    else
        update(mid+1,r,pos,val,2*rt+1);
    pushup(rt);
}

int query(int l,int r,int aa,int bb,int rt)
{
    if(aa <= l && bb >= r)
        return tree[rt];
    int mid = (l+r)/2;
    if(bb <= mid) return query(l,mid,aa,bb,2*rt);
    else if(aa > mid) return query(mid+1,r,aa,bb,2*rt+1);
    return max(query(l,mid,aa,bb,2*rt),query(mid+1,r,aa,bb,2*rt+1));
}

void dfs(int u,int f)
{
    dep[u] = dep[f]+1;
    siz[u] = 1;
    for(int i=head[u];i!=-1;i=G[i].next)
    {
        int v = G[i].v;
        if(v == f) continue;
        fa[v] = u;
        dfs(v,u);
        if(son[u] == -1 || siz[v] > siz[son[u]]) son[u] = v;
        siz[u] += siz[v];
    }
}

void dfs2(int u,int father)
{
    pos[u] = ++POS;
    Top[u] = father;
    if(son[u] != -1) dfs2(son[u],father);
    for(int i=head[u];i!=-1;i=G[i].next)
    {
        int v = G[i].v;
        if(v != fa[u] && v != son[u])
            dfs2(v,v);
    }
}

int ask(int u,int v)
{
    int fx = Top[u], fy = Top[v], maxi = 0;
    while(fx != fy)
    {
        if(dep[fx] < dep[fy])
        {
            swap(u,v);
            swap(fx,fy);
        }
        maxi = max(maxi,query(1,POS,pos[fx],pos[u],1));
        u = fa[fx];
        fx = Top[u];
    }
    if(u == v) return maxi;
    if(dep[u] > dep[v]) swap(u,v);
    return max(maxi,query(1,POS,pos[son[u]],pos[v],1));
}

int main()
{
    int u,v,w,x,y,i,t;
    char ss[10];
    scanf("%d",&t);
    while(t--)
    {
        scanf("%d",&n);
        init();
        for(i=1;i<n;i++)
        {
            scanf("%d%d%d",&edge[i].u,&edge[i].v,&edge[i].w);
            addedge(edge[i].u,edge[i].v);
        }
        dep[0] = 0;
        dfs(1,0);
        dfs2(1,1);
        for(i=1;i<n;i++)
        {
            if(dep[edge[i].u] > dep[edge[i].v])
                swap(edge[i].u,edge[i].v);
            update(1,POS,pos[edge[i].v],edge[i].w,1);
        }
        while(scanf("%s",ss)!=EOF && ss[0] != ‘D‘)
        {
            if(ss[0] == ‘Q‘)
            {
                scanf("%d%d",&u,&v);
                printf("%d\n",ask(u,v));
            }
            else
            {
                scanf("%d%d",&x,&y);
                update(1,POS,pos[edge[x].v],y,1);
            }
        }
        if(t >= 1)
            puts("");
    }
    return 0;
}