[ZJOI2008]树的统计 树链剖分

[ZJOI2008]树的统计

题面

裸树剖，一个线段树同时维护sum和max就好了。庆祝半小时1A此题。

另外注意，线段树同时支持「区间修改」和「区间查询」时才需要用到 lazy tag。

#include <cstdio>
#include <algorithm>
#define MAXN 30003
#define sl (x<<1)
#define sr (x<<1|1)
using namespace std;
int head[MAXN],nxt[MAXN*2],vv[MAXN*2],tot;
inline void add_edge(int u, int v){
	vv[++tot]=v;
	nxt[tot]=head[u];
	head[u]=tot;
}
int fa[MAXN],dep[MAXN],sz[MAXN],mxs[MAXN];
void dfs1(int u, int f){
	dep[u]=dep[f]+1;
	fa[u]=f;
	sz[u]=1;
	int mxsz=-1;
	for(int i=head[u];i;i=nxt[i]){
		int v=vv[i];
		if(v==f) continue;
		dfs1(v, u);
		sz[u]+=sz[v];
		if(mxsz<sz[v]){
			mxsz=sz[v];
			mxs[u]=v;
		}
	}
}
int w[MAXN],wnew[MAXN];
int topf[MAXN],idx[MAXN],cnt;
void dfs2(int u, int top){
	topf[u]=top;
	idx[u]=++cnt;
	wnew[cnt]=w[u];
	if(mxs[u]==0) return;
	dfs2(mxs[u], top);
	for(int i=head[u];i;i=nxt[i]){
		int v=vv[i];
		if(v==fa[u]||v==mxs[u]) continue;
		dfs2(v,v);
	}
}
struct nod{
	int sum,mx;
} tre[MAXN*4];
int n;
void push_up(int x){
	tre[x].sum=tre[sl].sum+tre[sr].sum;
	tre[x].mx=max(tre[sl].mx, tre[sr].mx);
}
void buildt(int x, int l, int r){
	if(l==r){
		tre[x].sum=tre[x].mx=wnew[l];
		return;
	}
	int mid=(l+r)>>1;
	buildt(sl, l, mid);
	buildt(sr, mid+1, r);
	push_up(x);
}
void change(int x, int cx, int l, int r, int val){
	if(l==r){
		tre[x].sum=tre[x].mx=val;
		return;
	}
	int mid=(l+r)>>1;
	if(cx<=mid) change(sl, cx, l, mid, val);
	else change(sr, cx, mid+1, r, val);
	push_up(x);
}
int query_mx(int x, int ql, int qr, int l, int r){
	if(ql<=l&&r<=qr){
		return tre[x].mx;
	}
	int mid=(l+r)>>1;
	int ans=-30000;
	if(ql<=mid) ans=max(query_mx(sl, ql, qr, l, mid), ans);
	if(mid<qr) ans=max(query_mx(sr, ql, qr, mid+1, r), ans);
	return ans;
}
int query_sum(int x, int ql, int qr, int l, int r){
	if(ql<=l&&r<=qr){
		return tre[x].sum;
	}
	int mid=(l+r)>>1;
	int ans=0;
	if(ql<=mid) ans+=query_sum(sl, ql, qr, l, mid);
	if(mid<qr) ans+=query_sum(sr, ql, qr, mid+1, r);
	return ans;
}
int tre_query_mx(int a, int b){
	int ans=-30000;
	while(topf[a]!=topf[b]){
		if(dep[topf[a]]<dep[topf[b]]) swap(a,b);
		ans=max(ans, query_mx(1, idx[topf[a]], idx[a], 1, n));
		a=fa[topf[a]];
	}
	if(dep[a]<dep[b]) swap(a,b);
	ans=max(query_mx(1, idx[b], idx[a], 1, n), ans);
	return ans;
}
int tre_query_sum(int a, int b){
	int ans=0;
	while(topf[a]!=topf[b]){
		if(dep[topf[a]]<dep[topf[b]]) swap(a,b);
		ans+=query_sum(1, idx[topf[a]], idx[a], 1, n);
		a=fa[topf[a]];
	}
	if(dep[a]<dep[b]) swap(a,b);
	ans+=query_sum(1, idx[b], idx[a], 1, n);
	return ans;
}
int main()
{
	scanf("%d", &n);
	for(int i=2;i<=n;++i){
		int a,b;
		scanf("%d %d", &a, &b);
		add_edge(a,b);
		add_edge(b,a);
	}
	for(int i=1;i<=n;++i) scanf("%d", &w[i]);
	dfs1(1, 1);
	dfs2(1, 1);
	buildt(1, 1, n);
	int q;
	scanf("%d", &q);
	while(q--){
		char s[10];
		int a,b;
		scanf("%s%d%d", s, &a, &b);
		if(s[1]=='H') change(1, idx[a], 1, n, b);
		else if(s[1]=='M') printf("%d\n", tre_query_mx(a,b));
		else if(s[1]=='S') printf("%d\n", tre_query_sum(a,b));
		else puts("Erro!");
	}
	return 0;
}