BZOJ 3531: [Sdoi2014]旅行(树链剖分+线段树)

传送门

解题思路

　　以每个颜色为根开一棵权值线段树，下标就是\(dfs\)序，其余都是基本操作，要动态开点。

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>

using namespace std;
const int N=100005;
const int M=N*20;

inline int rd(){
	int x=0,f=1;char ch=getchar();
	while(!isdigit(ch)) f=ch=='-'?0:1,ch=getchar();
	while(isdigit(ch)) x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
	return f?x:-x;
}	

inline int max(int x,int y) {return x>y?x:y;}

int n,q,head[N],cnt,to[N<<1],nxt[N<<1];
int C[N],W[N],siz[N],top[N];
int dep[N],id[N],fa[N],son[N],num,pool[N],cur;

struct Segment_tree{
	int rt[N],sum[M],ls[M],rs[M],tot,Max[M];
	inline int new_node(){
		return (cur?(pool[cur--]):(++tot));
	}
	void update(int &x,int l,int r,int pos,int k){
		if(!x) x=new_node(); 
		if(l==r) {
			sum[x]=k; Max[x]=k;
			if(!k) pool[++cur]=x,x=0;
			return;
		}
		int mid=(l+r)>>1;
		if(pos<=mid) update(ls[x],l,mid,pos,k);
		else update(rs[x],mid+1,r,pos,k);
		sum[x]=sum[ls[x]]+sum[rs[x]];
		Max[x]=max(Max[ls[x]],Max[rs[x]]);	
		if(!sum[x]) {pool[++cur]=x;x=0;}
	}
	int query_sum(int x,int l,int r,int L,int R){
		if(!x) return 0; if(L<=l && r<=R) return sum[x];
		int mid=(l+r)>>1,ret=0;
		if(L<=mid) ret+=query_sum(ls[x],l,mid,L,R);
		if(mid<R) ret+=query_sum(rs[x],mid+1,r,L,R);
		return ret;
	}
	int query_max(int x,int l,int r,int L,int R){
		if(!x) return 0; if(L<=l && r<=R) return Max[x];
		int mid=(l+r)>>1,ret=0;
		if(L<=mid) ret=max(ret,query_max(ls[x],l,mid,L,R));
		if(mid<R) ret=max(ret,query_max(rs[x],mid+1,r,L,R));
		return ret;
	}
	inline void BUILD(){
		for(int i=1;i<=n;i++) update(rt[C[i]],1,n,id[i],W[i]);
	}
}tree;

inline void add(int bg,int ed){
	to[++cnt]=ed,nxt[cnt]=head[bg],head[bg]=cnt;
}

void dfs1(int x,int f,int d){
	fa[x]=f; siz[x]=1; dep[x]=d;
	int maxson=-1,u;
	for(int i=head[x];i;i=nxt[i]){
		u=to[i]; if(u==f) continue;
		dfs1(u,x,d+1); siz[x]+=siz[u];
		if(siz[u]>maxson) maxson=siz[u],son[x]=u;
	}
}

void dfs2(int x,int topf){
	id[x]=++num; top[x]=topf; if(!son[x]) return;
	dfs2(son[x],topf); int u;
	for(int i=head[x];i;i=nxt[i]){
		u=to[i]; if(u==son[x] || u==fa[x]) continue;
		dfs2(u,u);
	}
}

inline int Qsum(int x,int y){
	int ret=0,CC=C[y];
	while(top[x]!=top[y]){
		if(dep[top[x]]<dep[top[y]]) swap(x,y);
		ret+=tree.query_sum(tree.rt[CC],1,n,id[top[x]],id[x]);
		x=fa[top[x]];
	}
	if(dep[x]>dep[y]) swap(x,y);
	ret+=tree.query_sum(tree.rt[CC],1,n,id[x],id[y]);
	return ret;
}

inline int Qmax(int x,int y){
	int ret=0,CC=C[y];
	while(top[x]!=top[y]){
		if(dep[top[x]]<dep[top[y]]) swap(x,y);
		ret=max(ret,tree.query_max(tree.rt[CC],1,n,id[top[x]],id[x]));
		x=fa[top[x]];
	}
	if(dep[x]>dep[y]) swap(x,y);
	ret=max(ret,tree.query_max(tree.rt[CC],1,n,id[x],id[y]));
	return ret;
}

int main(){
	n=rd(),q=rd(); int x,y; char s[5];
	for(int i=1;i<=n;i++) W[i]=rd(),C[i]=rd();
	for(int i=1;i<n;i++){
		x=rd(),y=rd();
		add(x,y); add(y,x);
	}
	dfs1(1,0,0); dfs2(1,1); tree.BUILD();
	while(q--){
		scanf("%s",s+1); x=rd(),y=rd();
		if(s[2]=='C') {
			tree.update(tree.rt[C[x]],1,n,id[x],0); C[x]=y;
			tree.update(tree.rt[C[x]],1,n,id[x],W[x]);
		} 
		else if(s[2]=='W') tree.update(tree.rt[C[x]],1,n,id[x],y),W[x]=y;
		else if(s[2]=='S') printf("%d\n",Qsum(x,y));
		else printf("%d\n",Qmax(x,y));
	}
	return 0;
}