POJ 2763 Housewife Wind（树链剖分+线段树）

#include <iostream>
#include <cstdio>
#include <cstring>
using namespace std;

const int maxn=500008;

int n,m,s;
int a[maxn],sz[maxn],dep[maxn],fa[maxn],top[maxn],son[maxn],head[maxn];
int w[maxn],num,cnt;
struct Edge
{
    int to,next;
};
Edge v[maxn*2];
void init()
{
    num=0;cnt=0;
    memset(head,-1,sizeof(head));
}
void addedge(int x,int y)
{
    v[cnt].to = y;
    v[cnt].next = head[x];
    head[x] = cnt++;
}
void dfs1(int u, int f, int d) {
    dep[u]=d;sz[u]=1;
    son[u]=0;fa[u]=f;
    for(int i=head[u];i!=-1;i = v[i].next){
        int to = v[i].to;
        if(to!=f){
            dfs1(to,u,d+1);
            sz[u]+=sz[to];
            if(sz[son[u]]<sz[to])son[u]=to;
        }
    }
}
void dfs2(int u, int tp) {
    top[u]=tp;
    w[u]=++num;
    if(son[u])dfs2(son[u],tp);
    for(int i = head[u];i!=-1;i=v[i].next){
        int to=v[i].to;
        if(to==fa[u]||to==son[u])continue;
        dfs2(to,to);
    }
}
struct Node
{
    int x,y,val;
}e[maxn];

int sum[maxn<<2];
void pushup(int rt)
{
    sum[rt]=sum[rt<<1]+sum[rt<<1|1];
}
void build(int l,int r,int rt)
{
    if(l==r){
        sum[rt]=a[l];
        return ;
    }
    int mid=(l+r)>>1;
    build(l,mid,rt<<1);
    build(mid+1,r,rt<<1|1);
    pushup(rt);
}
void update(int x,int val,int l,int r,int rt)
{
    if(l==r){
        sum[rt]=val;
        return ;
    }
    int mid=(l+r)>>1;
    if(x<=mid)update(x,val,l,mid,rt<<1);
    else update(x,val,mid+1,r,rt<<1|1);
    pushup(rt);
}
int query(int L,int R,int l,int r,int rt)
{
    if(L<=l&&r<=R){
        return sum[rt];
    }
    int mid=(l+r)>>1;
    int ans=0;
    if(L<=mid)ans+=query(L,R,l,mid,rt<<1);
    if(mid<R)ans+=query(L,R,mid+1,r,rt<<1|1);
    return ans;
}
int modify(int x,int y)
{
    int res=0;
    while(top[x]!=top[y]){
        if(dep[top[x]]<dep[top[y]])swap(x,y);
        res+=query(w[top[x]],w[x],1,num,1);
        x=fa[top[x]];
    }
    if(x==y)return res;
    if(dep[x]>dep[y])swap(x,y);
    res+=query(w[son[x]],w[y],1,num,1);
    return res;
}
int main()
{
    while(~scanf("%d%d%d",&n,&m,&s)){
        init();
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&e[i].x,&e[i].y,&e[i].val);
            addedge(e[i].x,e[i].y);
            addedge(e[i].y,e[i].x);
        }
        dfs1(1,0,1);
        dfs2(1,1);
        for(int i=1;i<n;i++){
            if(dep[e[i].x]<dep[e[i].y])
                swap(e[i].x,e[i].y);
            a[w[e[i].x]]=e[i].val;
        }
        build(1,num,1);
        for(int i=0;i<m;i++){
            int op,x,y;
            scanf("%d",&op);
            if(op==0){
                scanf("%d",&x);
                printf("%d\n",modify(s,x));
                s=x;
            }
            else{
                scanf("%d%d",&x,&y);
                update(w[e[x].x],y,1,num,1);
            }
        }
    }
    return 0;
}
/*
3 3 1
1 2 1
2 3 2
0 2
1 2 3
0 3

3 3 1
1 2 1
2 3 2
0 2
*/