HYSBZ - 1036 树的统计Count 树链剖分 求和+最大值

好水0.0

#include<cstdio>
#include<cstring>
#include<cmath>
#include<iostream>
#include<algorithm>
#include<set>
#include<map>
#include<queue>
#include<vector>
#include<string>
#define eps  1e-12
#define INF   0x7fffffff
#define maxn 31111
using namespace std;
char str[maxn];
struct node
{
    int to,next;
} e[maxn<<1];
int head[maxn],en,n;
int fa[maxn],siz[maxn],top[maxn],son[maxn],w[maxn],dep[maxn],id;
int maxs[maxn<<2];
int sums[maxn<<2];
int d[maxn][3];
void init()
{
    fa[1]=0;
    dep[1]=0;
    id=en=0;
    memset(head,-1,sizeof(head));
    memset(siz,0,sizeof(siz));
    memset(maxs,0,sizeof(maxs));
    memset(sums,0,sizeof(sums));
}
void add(int a,int b)
{
    e[en].to=b;
    e[en].next=head[a];
    head[a]=en++;
}
void dfs(int now)
{
    siz[now]=1;
    son[now]=0;
    for(int i=head[now]; ~i; i=e[i].next)
    {
        int to=e[i].to;
        if(to!=fa[now])
        {
            fa[to]=now;
            dep[to]=dep[now]+1;
            dfs(to);
            if(siz[to]>siz[son[now]]) son[now]=to;
            siz[now]+=siz[to];
        }
    }
}
void getid(int now,int root)
{
    w[now]=++id;
    top[now]=root;
    if(son[now]) getid(son[now],top[now]);
    for(int i=head[now]; ~i; i=e[i].next)
    {
        if(e[i].to!=son[now]&&e[i].to!=fa[now])
        {
            getid(e[i].to,e[i].to);
        }
    }
}
void update(int rt,int l,int r,int k,int add)
{
    if(l==r)
    {
        maxs[rt]=add;
        sums[rt]=add;
        return ;
    }
    int mid=(l+r)/2;
    int lson=(rt<<1);
    int rson=(rt<<1|1);
    if(k<=mid) update(lson,l,mid,k,add);
    else update(rson,mid+1,r,k,add);
    maxs[rt]=max(maxs[lson],maxs[rson]);
    sums[rt]=sums[lson]+sums[rson];
}
int querymax(int rt,int l,int r,int ql,int qr)
{
  //  printf("%d %d %d %d %d\n",l,r,ql,qr,maxs[rt]);
    if(ql<=l&&r<=qr) return maxs[rt];
    int mid=(l+r)/2;
    int lson=(rt<<1);
    int rson=(rt<<1|1);
    int tmp=-INF;
    if(ql<=mid) tmp=max(tmp,querymax(lson,l,mid,ql,qr));
    if(qr>mid)  tmp=max(tmp,querymax(rson,mid+1,r,ql,qr));
    return tmp;
}
int querysum(int rt,int l,int r,int ql,int qr)
{
    if(ql<=l&&r<=qr) return sums[rt];
    int mid=(l+r)/2;
    int lson=(rt<<1);
    int rson=(rt<<1|1);
    int tmp=0;
    if(ql<=mid) tmp+=querysum(lson,l,mid,ql,qr);
    if(qr>mid)  tmp+=querysum(rson,mid+1,r,ql,qr);
    return tmp;
}
int find(int a,int b,int which) //1求max 2求sum
{
    int f1=top[a];
    int f2=top[b];
    int ans=-INF,sum=0;
    while(f1!=f2)
    {
        if(dep[f1]>dep[f2])
        {
            if(which==1) ans=max(ans,querymax(1,1,n,w[f1],w[a]));
            else sum+=querysum(1,1,n,w[f1],w[a]);
            a=fa[f1];
            f1=top[a];
        }
        else
        {
            if(which==1) ans=max(ans,querymax(1,1,n,w[f2],w[b]));
            else sum+=querysum(1,1,n,w[f2],w[b]);
            b=fa[f2];
            f2=top[b];
        }
    }
    if(which==1){
        if(dep[a]>dep[b]) ans=max(ans,querymax(1,1,n,w[b],w[a]));
        else ans=max(ans,querymax(1,1,n,w[a],w[b]));
    }
    else
    {
        if(dep[a]>dep[b]) sum+=querysum(1,1,n,w[b],w[a]);
        else sum+=querysum(1,1,n,w[a],w[b]);
    }
    if(which==1) return ans;
    else return sum;
}
int main()
{
    char str[20];
    int a,b,c;
    while(~scanf("%d",&n))
    {
        init();
        for(int i=1;i<n;i++)
        {
            scanf("%d%d",&a,&b);
            add(a,b);
            add(b,a);
        }
        dfs(1);
        getid(1,1);
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&a);
            update(1,1,n,w[i],a);
        }

        scanf("%d",&c);
        while(c--)
        {
            scanf("%s%d%d",str,&a,&b);
            if(str[1]=='M')
            {
                printf("%d\n",find(a,b,1));
            }
            else if(str[1]=='S')
            {
                printf("%d\n",find(a,b,2));
            }
            else
            {
                update(1,1,n,w[a],b);
            }
        }
    }
    return 0;
}
/*
4
1 2
2 3
4 1
4 2 1 3
12
QMAX 3 4
QMAX 3 3
QMAX 3 2
QMAX 2 3
QSUM 3 4
QSUM 2 1
CHANGE 1 5
QMAX 3 4
CHANGE 3 6
QMAX 3 4
QMAX 2 4
QSUM 3 4
*/

HYSBZ - 1036 树的统计Count 树链剖分 求和+最大值