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

/**************************************************************
Problem: 1036
User: beginend
Language: C++
Result: Accepted
Time:2412 ms
Memory:4848 kb
****************************************************************/

#include <iostream>
#include <cstdio>
#include <cmath>

const int maxn=30005;

using namespace std;

int n,cnt,size[maxn],top[maxn],fa[maxn],dep[maxn],last[maxn],a[maxn],sz,pos[maxn];
//size为子树大小，top为链头，fa为父亲，dep为深度，pos为dfs序
struct data{int y,next;}edge[maxn*2];
struct tree{int mx,sum;}t[maxn*4];

void insert_edge(int u,int v)
{
edge[++cnt].y=v;edge[cnt].next=last[u];last[u]=cnt;
edge[++cnt].y=u;edge[cnt].next=last[v];last[v]=cnt;
}

void initt()
{
int i,u,v;
scanf("%d",&n);
for (i=1;i<=n-1;i++)
{
scanf("%d%d",&u,&v);
insert_edge(u,v);
}
for (i=1;i<=n;i++)
scanf("%d",&a[i]);
}

void dfs1(int x)
{
int c=last[x];
size[x]=1;
while (c>0)
{
int y=edge[c].y;
if (y!=fa[x])
{
fa[y]=x;
dep[y]=dep[x]+1;
dfs1(y);
size[x]+=size[y];
}
c=edge[c].next;

}
}

void dfs2(int x,int ch)
{
int k=0; sz++;
pos[x]=sz;
top[x]=ch;
for (int c=last[x];c>0;c=edge[c].next)
{
int y=edge[c].y;
if ((fa[y]==x) && (size[y]>size[k]))
{
k=y;
}
}
if (k==0) return;
dfs2(k,ch);
for (int c=last[x];c>0;c=edge[c].next)
{
int y=edge[c].y;
if ((fa[y]==x) && (y!=k)) dfs2(y,y);
}
}

void change(int p,int l,int r,int x,int c)
{
if (l==r)
{
t[p].mx=c;
t[p].sum=c;
return;
}
int mid=(l+r)/2;
if (x<=mid) change(p*2,l,mid,x,c);
else change(p*2+1,mid+1,r,x,c);
t[p].sum=t[p*2].sum+t[p*2+1].sum;
t[p].mx=max(t[p*2].mx,t[p*2+1].mx);
}

int getmax(int p,int l,int r,int x,int y)
{
if ((l==x) && (r==y))
{
return t[p].mx;
}
int mid=(l+r)/2;
if (y<=mid) return getmax(p*2,l,mid,x,y);
else
{
if (x>mid)
{
return getmax(p*2+1,mid+1,r,x,y);
}
else return max(getmax(p*2,l,mid,x,mid),getmax(p*2+1,mid+1,r,mid+1,y));
}
}

int getsum(int p,int l,int r,int x,int y)
{
if ((l==x) && (r==y))
{
return t[p].sum;
}
int mid=(l+r)/2;
if (y<=mid) return getsum(p*2,l,mid,x,y);
else
{
if (x>mid)
{
return getsum(p*2+1,mid+1,r,x,y);
}
else return getsum(p*2,l,mid,x,mid)+getsum(p*2+1,mid+1,r,mid+1,y);
}
}

int ask_max(int x,int y)
{
int ans=-0x3f3f3f3f;
while (top[x]!=top[y])
{
if (dep[top[x]]>dep[top[y]]) swap(x,y);
ans=max(ans,getmax(1,1,n,pos[top[y]],pos[y]));
y=fa[top[y]];
}
if (dep[x]>dep[y]) swap(x,y);
ans=max(ans,getmax(1,1,n,pos[x],pos[y]));
return ans;
}

int ask_sum(int x,int y)
{
int ans=0;
while (top[x]!=top[y])
{
if (dep[top[x]]>dep[top[y]]) swap(x,y);
ans+=getsum(1,1,n,pos[top[y]],pos[y]);
y=fa[top[y]];
}
if (dep[x]>dep[y]) swap(x,y);
ans+=getsum(1,1,n,pos[x],pos[y]);
return ans;
}

int main()
{
initt();
dfs1(1);
dfs2(1,1);
for (int i=1;i<=n;i++)
change(1,1,n,pos[i],a[i]);
int test;
scanf("%d",&test);
char ch[10]; int x,y;
for (int i=1;i<=test;i++)
{
scanf("%s",ch);
scanf("%d%d",&x,&y);
if (ch[1]=='H') change(1,1,n,pos[x],y);
if (ch[1]=='M') printf("%d\n",ask_max(x,y));
if (ch[1]=='S') printf("%d\n",ask_sum(x,y));
}
return 0;
}