bzoj 1036 [ZJOI2008]树的统计Count 线段树+树链剖分

Code

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define N 30010
#define inf 0x3f3f3f3f
#define ls rt<<1
#define rs rt<<1|1
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
using namespace std;
struct ss
{
int next,to;
};
ss Edge[N<<1];
int n,m,tot,QT;
int f[N<<2];
int a
;
int w
;
int fa
;
int sum[N<<2];
int top
;
int ran
;
int dfn
;
int son
;
int head
;
int size
;

void addedge(int x,int y)
{
Edge[++tot].next=head[x];
Edge[tot].to=y;
head[x]=tot;
}

void dfs1(int u,int root,int deep)
{
dfn[u]=deep;
fa[u]=root;
size[u]=1;
for(int i=head[u];i;i=Edge[i].next)
{
int to=Edge[i].to;
if(to==root)continue;
dfs1(to,u,deep+1);
size[u]+=size[to];
if(size[son[u]]<size[to])son[u]=to;
}
}

void dfs2(int u,int tp)
{
top[u]=tp;
w[u]=++QT;
ran[QT]=u;
if(!son[u])return ;
dfs2(son[u],tp);
for(int i=head[u];i;i=Edge[i].next)
{
int to=Edge[i].to;
if(to==fa[u]||to==son[u])continue;
dfs2(to,to);
}
}

void PushUp(int rt)
{
f[rt]=max(f[ls],f[rs]);
sum[rt]=sum[ls]+sum[rs];
}

void build(int l,int r,int rt)
{
if(l==r)
{
f[rt]=sum[rt]=a[ran[l]];
return ;
}
int mid=(l+r)>>1;
build(lson);
build(rson);
PushUp(rt);
}

void updata(int l,int r,int rt,int x,int c)
{
if(l==r)
{
f[rt]=sum[rt]=c;
return ;
}
int mid=(l+r)>>1;
if(mid>=x)updata(lson,x,c);
else updata(rson,x,c);
PushUp(rt);
}

int query_max(int l,int r,int rt,int L,int R)
{
if(l>=L&&r<=R)return f[rt];
int mid=(l+r)>>1,ans=-inf;
if(mid>=L)ans=max(ans,query_max(lson,L,R));
if(mid<R)ans=max(ans,query_max(rson,L,R));
return ans;
}

int ask_max(int x,int y)
{
int ans=-inf;
while(top[x]!=top[y])
{
if(dfn[top[x]]<dfn[top[y]])swap(x,y);
ans=max(query_max(1,n,1,w[top[x]],w[x]),ans);
x=fa[top[x]];
}
if(dfn[x]>dfn[y])swap(x,y);
ans=max(query_max(1,n,1,w[x],w[y]),ans);
return ans;
}

int query_sum(int l,int r,int rt,int L,int R)
{
if(l>=L&&r<=R)return sum[rt];
int mid=(l+r)>>1;
int ans=0;
if(mid>=L)ans+=query_sum(lson,L,R);
if(mid<R)ans+=query_sum(rson,L,R);
return ans;
}

int ask_sum(int x,int y)
{
int ans=0;
while(top[x]!=top[y])
{
if(dfn[top[x]]<dfn[top[y]])swap(x,y);
ans+=query_sum(1,n,1,w[top[x]],w[x]);
x=fa[top[x]];
}
if(dfn[x]>dfn[y])swap(x,y);
ans+=query_sum(1,n,1,w[x],w[y]);
return ans;
}

int main()
{
cin>>n;
for(int x,y,i=1;i<n;i++)
{
scanf("%d%d",&x,&y);
addedge(x,y),addedge(y,x);
}
for(int i=1;i<=n;i++)scanf("%d",&a[i]);
dfs1(1,0,1);
dfs2(1,1);
build(1,n,1);
cin>>m;
for(int i=1;i<=m;i++)
{
char s[10];
int x,y;
scanf("%s%d%d",s,&x,&y);
if(s[0]=='C')updata(1,n,1,w[x],y);
else{
if(s[1]=='M')printf("%d\n",ask_max(x,y));
else printf("%d\n",ask_sum(x,y));
}
}
return 0;
}