Bzoj 1036: [ZJOI2008]树的统计Count

题目链接：

http://www.lydsy.com/JudgeOnline/problem.php?id=1036

题解：

树链剖分＋线段树操作

代码：

#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
#define met(a,b) memset(a,b,sizeof(a))
#define inf 0x7fffffff
const int maxn = 30005;
struct EDGE
{
    int to,next;
}Edge[maxn*2];
struct Tree
{
    int left,right;
    int sum,MAX;
}tree[maxn<<2];
int head[maxn];
int cnt;
void init1()
{
    cnt=0;
    met(head,-1);
}

void Add_edge(int u,int v)
{
    Edge[++cnt].to=v,Edge[cnt].next=head[u],head[u]=cnt;
    Edge[++cnt].to=u,Edge[cnt].next=head[v],head[v]=cnt;
}

int size[maxn],fa[maxn],deep[maxn];
void init2()
{
    met(size,0),met(fa,0),met(deep,0);
}

//size[i]:表示i这个节点下面的字节点的个数

void Dfs1(int x)
{
    size[x]=1;
    for(int i=head[x];i!=-1;i=Edge[i].next)
    {
        if(Edge[i].to==fa[x])
            continue;
        deep[Edge[i].to]=deep[x]+1;
        fa[Edge[i].to]=x;
        Dfs1(Edge[i].to);
        size[x]+=size[Edge[i].to];
    }
}

int bl[maxn],pos[maxn];
int sz;
void init3()
{
    sz=0;
    met(bl,0),met(pos,0);
}

//pos[i]:表示i这个节点在线段树的位置
void Dfs2(int x,int chain)
{
    int k=0;
    sz++;
    pos[x]=sz;//分配在线段树里的位置
    bl[x]=chain;
    for(int i=head[x];i!=-1;i=Edge[i].next)
    {
        if(deep[Edge[i].to]>deep[x]&&size[Edge[i].to]>size[k])
            k=Edge[i].to;//选择子树中的最大的儿子来继承重链
    }
    if(k==0)
        return;
    Dfs2(k,chain);
    for(int i=head[x];i!=-1;i=Edge[i].next)
    {
        if(deep[Edge[i].to]>deep[x]&&k!=Edge[i].to)
            Dfs2(Edge[i].to,Edge[i].to);
    }
}

//－－－－－－－线段树
void build(int rt,int l,int r)
{
    tree[rt].left=l;
    tree[rt].right=r;
    if(l==r)return;
    int mid=(l+r)>>1;
    build(rt<<1,l,mid);
    build(rt<<1|1,mid+1,r);
}

void updata(int rt,int pos,int delta)
{
    int left=tree[rt].left,right=tree[rt].right;
    int mid=(left+right)>>1;
    if(left==right)
    {
        tree[rt].sum=tree[rt].MAX=delta;
        return;
    }
    if(pos<=mid)
        updata(rt<<1,pos,delta);
    else
        updata(rt<<1|1,pos,delta);
    tree[rt].sum=tree[rt<<1].sum+tree[rt<<1|1].sum;
    tree[rt].MAX=max(tree[rt<<1].MAX,tree[rt<<1|1].MAX);
}

int querysum(int x,int y,int rt)
{
    int left=tree[rt].left,right=tree[rt].right,mid=(left+right)>>1;
    if(left==x&&right==y)
        return tree[rt].sum;
    if(y<=mid)
        return querysum(x,y,rt<<1);
    else if(x>mid)
        return querysum(x,y,rt<<1|1);
    else
        return querysum(x,mid,rt<<1)+querysum(mid+1,y,rt<<1|1);
}

int querymax(int x,int y,int rt)
{
    int left=tree[rt].left,right=tree[rt].right,mid=(left+right)>>1;
    if(left==x&&right==y)
        return tree[rt].MAX;
    if(y<=mid)
        return querymax(x,y,rt<<1);
    else if(x>mid)
        return querymax(x,y,rt<<1|1);
    else
        return max(querymax(x,mid,rt<<1),querymax(mid+1,y,rt<<1|1));
}

int slovemax(int x,int y)
{
    int MAX=-inf;
    while(bl[x]!=bl[y])
    {
        if(deep[bl[x]]<deep[bl[y]])
            swap(x,y);
        MAX=max(MAX,querymax(pos[bl[x]],pos[x],1));
        x=fa[bl[x]];
    }
    if(pos[x]>pos[y])
        swap(x,y);
    MAX=max(MAX,querymax(pos[x],pos[y],1));
    return MAX;
}

int slovesum(int x,int y)
{
    int sum=0;
    while(bl[x]!=bl[y])
    {
        if(deep[bl[x]]<deep[bl[y]])
            swap(x,y);
        sum+=querysum(pos[bl[x]],pos[x],1);
        x=fa[bl[x]];
    }
    if(pos[x]>pos[y])
        swap(x,y);
    sum+=querysum(pos[x],pos[y],1);
    return sum;
}

int room[maxn];
int main()
{
    int n;
    scanf("%d",&n);
    init1(),init2(),init3();
    for(int i=1;i<n;i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        Add_edge(x,y);
    }
    met(room,0);
    for(int i=1;i<=n;i++)
        scanf("%d",&room[i]);
    Dfs1(1);
    Dfs2(1,1);
    build(1,1,n);
    for(int i=1;i<=n;i++)
        updata(1,pos[i],room[i]);
    int m;
    scanf("%d",&m);
    char s[maxn];
    int x,y;
    while(m--)
    {
        scanf("%s%d%d",s,&x,&y);
        if(strcmp(s,"CHANGE")==0)
            updata(1,pos[x],y);
        else if(strcmp(s,"QMAX")==0)
            printf("%d\n",slovemax(x,y));
        else
            printf("%d\n",slovesum(x,y));
    }
    return 0;
}