[BZOJ] [P2243] [SDOI 2011] [염색] [문제 풀이] [나무 사슬 분할]

#include<cstdio>
#include<vector>
#include<iostream>
#include<algorithm>
#include<climits>
#define lson i<<1,l,mid
#define rson i<<1|1,mid+1,r
#define L i<<1
#define R i<<1|1
using namespace std;
const int maxn=1e5+5;
int top[maxn],fa[maxn],dep[maxn],son[maxn],siz[maxn],w[maxn],z=0;
vector<int>G[maxn];

int n,m,root=1;
int c[maxn],col[maxn];
int U[maxn],V[maxn];
void add(int u,int v){
	G[u].push_back(v);
	G[v].push_back(u);
}	
struct node{
	int l,r,sum,lazy;
	node(){
		l=-1;r=-1;sum=0,lazy=-1;
	}
	node(int _l,int _r,int _sum){
		l=_l;r=_r;sum=_sum;
	}
};
node operator+(node a,node b){
	if(a.sum==0)return b;
	if(b.sum==0)return a;
	node res;
	res.l=a.l;
	res.r=b.r;
	res.sum=a.sum+b.sum-(a.r==b.l);		
	return res;
}
struct seg_tree{

	node t[maxn<<2];
	/*void build(int i,int l,int r){
		if(l>r)return;
		if(l==r){
			t[i].l=t[i].r=c[l];
			t[i].sum=1;
		}
		int mid=(l+r)>>1;
		build(lson);build(rson);
		t[i]=t[L]+t[R];
	}*/
	void pushdown(int i){
		if(t[i].lazy==-1)return;
		t[L].l=t[L].r=t[R].l=t[R].r=t[i].lazy;
		t[L].sum=t[R].sum=1;
		t[L].lazy=t[R].lazy=t[i].lazy;
		t[i].lazy=-1;
	}
	void updata(int i,int l,int r,int l0,int r0,int c){
		if(l>r)return;
		if(l0<=l&&r0>=r){
			t[i].l=t[i].r=c;
			t[i].sum=1;
			t[i].lazy=c;
			//if(l0==l&&r0==r)t[i].lazy=-1;
			return;
		}
		int mid=(l+r)>>1;
		pushdown(i);
		if(l0<=mid)updata(lson,l0,r0,c);
		if(r0>mid)updata(rson,l0,r0,c);
		t[i].l=t[L].l;
		t[i].r=t[R].r;
		t[i].sum=t[L].sum+t[R].sum-(t[L].r==t[R].l&&t[L].r!=-1);
	} 
	node query(int i,int l,int r,int l0,int r0){
		if(l>r)return node(-1,-1,0);
		if(l0<=l&&r0>=r)return t[i];
		int mid=l+r>>1;
		pushdown(i);
		if(r0<=mid)return query(lson,l0,r0);
		if(l0>mid)return query(rson,l0,r0);
		return query(lson,l0,mid)+query(rson,mid+1,r0);
	}
	
}T;
void dfs(int u){
	son[u]=0;siz[u]=1;
	for(int i=0;i<G[u].size();i++){
		int v=G[u][i];
		if(v!=fa[u]){
			fa[v]=u;
			dep[v]=dep[u]+1;
			dfs(v);
			if(siz[v]>siz[son[u]])son[u]=v;
			siz[u]+=siz[v];
		}
	}
}
void build(int u,int tp){
	w[u]=++z;top[u]=tp;
	if(son[u]!=0)build(son[u],tp);
	for(int i=0;i<G[u].size();i++){
		int v=G[u][i];
		if(v!=son[u]&&v!=fa[u])
		build(v,v);
	}
}
void change(int u,int v,int c){
	while(top[u]!=top[v]){
		if(dep[top[u]]>dep[top[v]]){
			int a=w[u],b=w[top[u]];
			if(a>b)swap(a,b);
			T.updata(1,1,n,a,b,c);
			u=fa[top[u]];
		}else{
			int a=w[v],b=w[top[v]];
			if(a>b)swap(a,b);
			T.updata(1,1,n,a,b,c);
			v=fa[top[v]];
		}
	}
	if(w[u]<w[v])
		T.updata(1,1,n,w[u],w[v],c);
	else
		T.updata(1,1,n,w[v],w[u],c);
}
int query(int u,int v){
	node Left,Right,ans;
	while(top[u]!=top[v]){
		if(dep[top[u]]>dep[top[v]]){
			int a=w[u],b=w[top[u]];
			ans=T.query(1,1,n,min(a,b),max(a,b));
			if(a>b)swap(ans.l,ans.r);
			Left=Left+ans;
			u=fa[top[u]];
		}else{
			int a=w[v],b=w[top[v]];
			ans=T.query(1,1,n,min(a,b),max(a,b));
			if(a<b)swap(ans.l,ans.r);
			Right=ans+Right;
			v=fa[top[v]];
		}
	}
	
	ans=T.query(1,1,n,min(w[u],w[v]),max(w[u],w[v]));
	
	if(w[u]>w[v])swap(ans.l,ans.r);
	Left=Left+ans;
	
	return (Left+Right).sum;
}
void deb(){	
	//DEB
	for(int i=1;i<=14;i++)
	printf("#%d l:%d r:%d sum:%d lazy:%d
",i,T.t[i].l,T.t[i].r,T.t[i].sum,T.t[i].lazy);
	//end DEB	
}
int main(){
	freopen("1.txt","r",stdin);
	freopen("3.txt","w",stdout);
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++)scanf("%d",&c[i]);
	for(int i=1;i<n;i++){
		int u,v;scanf("%d%d",&u,&v);
		add(u,v);U[i]=u;V[i]=v;
	}
	dep[root]=1;
	dfs(root);build(root,root);

	for(int i=1;i<=n;i++)
		change(i,i,c[i]);//,deb();



	
	
	
	char opt[5];
	while(m--){
		scanf("%s",opt);
		if(opt[0]=='Q'){
			int u,v;
			scanf("%d%d",&u,&v);
			printf("%d
",query(u,v));
		}else{
			int u,v,C;
			scanf("%d%d%d",&u,&v,&C);
			change(u,v,C);
		}//deb();
	}
	return 0;
}

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstdlib>
#include<ctime>
using namespace  std;
int main(){
	freopen("1.txt","w",stdout);
	srand(time(0));
	int n=10,m=10;
	cout<<n<<" "<<m<<endl;
	for(int i=1;i<=n;i++)
	cout<<rand()%3<<" ";
	cout<<endl;
	for(int i=2;i<=n;i++){
		int u=i;
		int v=rand()%u+1;
		if(u==v){
			i--;
			continue;
		}
		cout<<u<<" "<<v<<endl;
	}
	while(m--){
		if(rand()%2){
			int a=rand()%n+1,b=rand()%n+1,c=rand()%3;
			cout<<'C';
			cout<<" "<<a<<" "<<b<<" "<<c<<endl;
		}
		else{
			cout<<'Q';
			cout<<" "<<rand()%n+1<<" "<<rand()%n+1<<endl;
		}
	}
	return 0;
}

@echo off
:loop
pai.exe
2.exe
3.exe
fc 2.txt 3.txt
if %errorlevel%==0 goto loop
pause