bzoj1036 树的统计 树链剖分模板

#include<bits/stdc++.h>
using namespace std;
const int N = 1e5+5;
struct Edge
{
    int v,next;
}edge[N<<1];
int sum[N<<2],mx[N<<2],n;
int head[N],tot,dep[N],fa[N],sz[N],son[N],top[N],id[N],rk[N],cnt,val[N];
void init()
{
    memset(head,-1, sizeof(head));
    tot=0;
}
void add(int u,int v)
{
    edge[tot].v=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
void dfs1(int u,int f)
{
    dep[u]=dep[f]+1;
    fa[u]=f;
    sz[u]=1;
    for(int i=head[u];~i;i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==f)continue;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]])son[u]=v;
    }
}
void dfs2(int u,int t)
{
    top[u]=t;
    id[u]=++cnt;
    rk[cnt]=u;
    if(!son[u])return;
    dfs2(son[u],t);
    for(int i=head[u];~i;i=edge[i].next)
    {
        int v=edge[i].v;
        if(v!=son[u]&&v!=fa[u])dfs2(v,v);
    }
}
void pushup(int rt)
{
    sum[rt]=sum[rt<<1]+sum[rt<<1|1];
    mx[rt]=max(mx[rt<<1],mx[rt<<1|1]);
}
void build(int l,int r,int rt)
{
    if(l==r)
    {
        mx[rt]=sum[rt]=val[rk[l]];
        return;
    }
    int m=(l+r)>>1;
    build(l,m,rt<<1);
    build(m+1,r,rt<<1|1);
    pushup(rt);
}
int querySum(int L,int R,int l,int r,int rt)
{
    if(L<=l&&r<=R)return sum[rt];
    int m=(l+r)>>1;
    int res=0;
    if(L<=m)res+=querySum(L,R,l,m,rt<<1);
    if(R>m)res+=querySum(L,R,m+1,r,rt<<1|1);
    return res;
}
int queryMax(int L,int R,int l,int r,int rt)
{
    if(L<=l&&r<=R)return mx[rt];
    int m=(l+r)>>1;
    int res=-1e9;
    if(L<=m)res=max(res,queryMax(L,R,l,m,rt<<1));
    if(R>m)res=max(res,queryMax(L,R,m+1,r,rt<<1|1));
    return res;
}
void update(int pos,int val,int l,int r,int rt)
{
    if(l==r){
        sum[rt]=mx[rt]=val;
        return;
    }
    int m=(l+r)>>1;
    if(pos<=m)update(pos,val,l,m,rt<<1);
    else update(pos,val,m+1,r,rt<<1|1);
    pushup(rt);
}
int getSum(int x,int y)
{
    int ans=0;
    while(top[x]!=top[y])
    {
        if(dep[top[x]]<dep[top[y]])swap(x,y);
        ans+=querySum(id[top[x]],id[x],1,n,1);
        x=fa[top[x]];
    }
    if(id[x]>id[y])swap(x,y);
    ans+=querySum(id[x],id[y],1,n,1);
    return ans;
}
int getMax(int x,int y)
{
    int ans=-1e9;
    while (top[x]!=top[y])
    {
        if(dep[top[x]]<dep[top[y]])swap(x,y);
        ans=max(ans,queryMax(id[top[x]],id[x],1,n,1));
        x=fa[top[x]];
    }
    if(id[x]>id[y])swap(x,y);
    ans=max(ans,queryMax(id[x],id[y],1,n,1));
    return ans;
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    init();
    int u,v;
    cin>>n;
    for(int i=1;i<=n-1;i++)
    {
        cin>>u>>v;
        add(u,v);
        add(v,u);
    }
    for(int i=1;i<=n;i++)cin>>val[i];
    dfs1(1,1);
    dfs2(1,1);
    build(1,n,1);
    int q,x,y;
    string op;
    cin>>q;
    while(q--)
    {
        cin>>op>>x>>y;
        if(op=="QMAX")cout<<getMax(x,y)<<'\n';
        else if(op=="QSUM")cout<<getSum(x,y)<<'\n';
        else update(id[x],y,1,n,1);
    }
    return 0;
}