P2590 [ZJOI2008]树的统计题解

重链剖分

#include<bits/stdc++.h>
using namespace std;
const int maxn=3e4+10;
struct node{
	int to;
	int nxt;
}E[maxn<<1];
int n,h[maxn],etot,tot,top[maxn],siz[maxn],rnk[maxn],son[maxn],dep[maxn],fa[maxn],w[maxn],dfn[maxn];
void addedge(int x,int y)
{
	E[++etot].to=y;
	E[etot].nxt=h[x];
	h[x]=etot;	
}
//线段树模版
struct Segment{
	int sum;
	int max;
}Seg[maxn<<2];
void pushup(int rt)
{
	Seg[rt].max=max(Seg[rt<<1].max,Seg[rt<<1|1].max);
	Seg[rt].sum=Seg[rt<<1].sum+Seg[rt<<1|1].sum;
}
void build(int rt,int l,int r)
{
	if (l==r)
	{
		Seg[rt].max=Seg[rt].sum=w[rnk[l]];
		return ;
	}
	int mid=(l+r)>>1;
	build(rt<<1,l,mid);
	build(rt<<1|1,mid+1,r);
	pushup(rt);
}
void update(int rt,int l,int r,int ql,int val)
{
	if (l==r)
	{
		Seg[rt].max=Seg[rt].sum=val;
		return ;
	}
	int mid=(l+r)>>1;
	if (ql<=mid) update(rt<<1,l,mid,ql,val);
	else update(rt<<1|1,mid+1,r,ql,val);
	pushup(rt);
	
}
int querysum(int rt,int l,int r,int ql,int qr)
{
	if (ql<=l&&r<=qr)
	{
		return Seg[rt].sum;
	}
	int mid=(l+r)>>1,ans=0;
	if (ql<=mid) ans+=querysum(rt<<1,l,mid,ql,qr);
	if (mid<qr) ans+=querysum(rt<<1|1,mid+1,r,ql,qr);
	return ans;
}
int querymax(int rt,int l,int r,int ql,int qr)
{
	if (ql<=l&&r<=qr)
	{
		return Seg[rt].max;
	}
	int mid=(l+r)>>1,ans=0x80000000;
	if (ql<=mid) ans=max(ans,querymax(rt<<1,l,mid,ql,qr));
	if (mid<qr) ans=max(ans,querymax(rt<<1|1,mid+1,r,ql,qr));
	return ans;
}
//树链剖分重链部分两个dfs
void dfs1(int u,int f)
{
	siz[u]=1;
	fa[u]=f;
	dep[u]=dep[f]+1;
	for (int i=h[u];i;i=E[i].nxt)
	{
		int v=E[i].to;
		if (v!=f)
		{
			dfs1(v,u);
			siz[u]+=siz[v];
			if (siz[v]>siz[son[u]]) son[u]=v;
		}
	}
}
void dfs2(int u,int t)
{
	dfn[u]=++tot;
	rnk[tot]=u;
	top[u]=t;
	if (son[u]) dfs2(son[u],t);
	for (int i=h[u];i;i=E[i].nxt)
	{
		int v=E[i].to;
		if (v!=fa[u]&&v!=son[u]) dfs2(v,v); 
	}
}
//在若干个重链进行查询
int Qmax(int x,int y)
{
	int res=0x80000000;
	while(top[x]!=top[y])
	{
		if (dep[top[x]]<dep[top[y]]) swap(x,y);
		res=max(res,querymax(1,1,n,dfn[top[x]],dfn[x]));
		x=fa[top[x]];
	}
	if (dep[x]<dep[y]) swap(x,y);
	res=max(res,querymax(1,1,n,dfn[y],dfn[x]));
	return res;
}
int Qsum(int x,int y)
{
	int res=0;
	while(top[x]!=top[y])
	{
		if (dep[top[x]]<dep[top[y]]) swap(x,y);
		res+=querysum(1,1,n,dfn[top[x]],dfn[x]);
		x=fa[top[x]];
	}
	if (dep[x]<dep[y]) swap(x,y);
	res+=querysum(1,1,n,dfn[y],dfn[x]);
	return res; 
}
int main()
{
	int q,u,v;
	cin>>n;
	for (int i=1;i<n;i++)
	{
		int a,b;
		cin>>a>>b;
		addedge(a,b);
		addedge(b,a);
	}
	for (int i=1;i<=n;i++)
	{
		cin>>w[i];
	}
	dfs1(1,0);
	dfs2(1,1);
	build(1,1,n);
	cin>>q;
	for (int i=1;i<=q;i++)
	{
		string opt;
		int u,v;
		cin>>opt;
		if(opt=="CHANGE")
		{
			cin>>u>>v;
			update(1,1,n,dfn[u],v);
		}
		else if(opt=="QMAX")
		{
			cin>>u>>v;
			cout<<Qmax(u,v)<<endl;
		}
		else
		{
			cin>>u>>v;
			cout<<Qsum(u,v)<<endl;
		}
	}
}