luogu2486 [SDOI2011]染色

维护区间左颜色值，右颜色值，颜色段个数。

#include <iostream>
#include <cstdio>
using namespace std;
int n, m, uu, vv, ww, w[100005], wt[100005], qwq, idx[100005], hea[100005];
int cnt, dep[100005], siz[100005], son[100005], fa[100005], top[100005];
char ss[15];
struct Edge{
	int too, nxt;
}edge[200005];
struct SGT{
	int col[400005], lco[400005], rco[400005], tag[400005];
	void upd(int o, int lson, int rson){
		col[o] = col[lson] + col[rson];
		if(rco[lson]==lco[rson])	col[o]--;
		lco[o] = lco[lson];
		rco[o] = rco[rson];
	}
	void build(int o, int l, int r){
		if(l==r)	lco[o] = rco[o] = wt[l], col[o] = 1;
		else{
			int mid=(l+r)>>1;
			int lson=o<<1;
			int rson=lson|1;
			if(l<=mid)	build(lson, l, mid);
			if(mid<r)	build(rson, mid+1, r);
			upd(o, lson, rson);
		}
	}
	void pushDown(int o, int l, int r, int lson, int rson, int mid){
		col[lson] = col[rson] = 1;
		lco[lson] = rco[lson] = lco[rson] = rco[rson] = tag[o];
		tag[lson] = tag[rson] = tag[o];
		tag[o] = 0;
	}
	void update(int o, int l, int r, int x, int y, int k){
		if(l>=x && r<=y){
			col[o] = 1;
			lco[o] = rco[o] = tag[o] = k;
		}
		else{
			int mid=(l+r)>>1;
			int lson=o<<1;
			int rson=lson|1;
			if(tag[o])	pushDown(o, l, r, lson, rson, mid);
			if(x<=mid)	update(lson, l, mid, x, y, k);
			if(mid<y)	update(rson, mid+1, r, x, y, k);
			upd(o, lson, rson);
		}
	}
	int query(int o, int l, int r, int x, int y){
		if(l>=x && r<=y)	return col[o];
		else{
			int mid=(l+r)>>1;
			int lson=o<<1;
			int rson=lson|1;
			int ans=0;
			if(tag[o])	pushDown(o, l, r, lson, rson, mid);
			if(x<=mid)	ans += query(lson, l, mid, x, y);
			if(mid<y)	ans += query(rson, mid+1, r, x, y);
			if(x<=mid && mid<y && rco[lson]==lco[rson])	return ans-1;
			return ans;
		}
	}
	int queryCol(int o, int l, int r, int x){
		if(l==r)	return lco[o];
		else{
			int mid=(l+r)>>1;
			int lson=o<<1;
			int rson=lson|1;
			if(tag[o])	pushDown(o, l, r, lson, rson, mid);
			if(x<=mid)	return queryCol(lson, l, mid, x);
			else	return queryCol(rson, mid+1, r, x);
		}
	}
}sgt;
void add_edge(int fro, int too){
	edge[++cnt].nxt = hea[fro];
	edge[cnt].too = too;
	hea[fro] = cnt;
}
void dfs1(int x, int f){
	dep[x] = dep[f] + 1;
	siz[x] = 1;
	fa[x] = f;
	int maxSon=-1;
	for(int i=hea[x]; i; i=edge[i].nxt){
		int t=edge[i].too;
		if(t!=f){
			dfs1(t, x);
			siz[x] += siz[t];
			if(maxSon<siz[t])	maxSon = siz[t], son[x] = t;
		}
	}
}
void dfs2(int x, int topf){
	top[x] = topf;
	idx[x] = ++qwq;
	wt[qwq] = w[x];
	if(!son[x])	return ;
	dfs2(son[x], topf);
	for(int i=hea[x]; i; i=edge[i].nxt){
		int t=edge[i].too;
		if(t!=fa[x] && t!=son[x])
			dfs2(t, t);
	}
}
void changeRange(int uu, int vv, int ww){
	while(top[uu]!=top[vv]){
		if(dep[top[uu]]<dep[top[vv]])	swap(uu, vv);
		sgt.update(1, 1, n, idx[top[uu]], idx[uu], ww);
		uu = fa[top[uu]];
	}
	if(dep[uu]>dep[vv])	swap(uu, vv);
	sgt.update(1, 1, n, idx[uu], idx[vv], ww);
}
int queryRange(int uu, int vv){
	int ans=0, orz, sto;
	while(top[uu]!=top[vv]){
		if(dep[top[uu]]<dep[top[vv]])	swap(uu, vv);
		ans += sgt.query(1, 1, n, idx[top[uu]], idx[uu]);
		orz = sgt.queryCol(1, 1, n, idx[top[uu]]);
		sto = sgt.queryCol(1, 1, n, idx[fa[top[uu]]]);
		if(orz==sto)	ans--;
		uu = fa[top[uu]];
	}
	if(dep[uu]>dep[vv])	swap(uu, vv);
	ans += sgt.query(1, 1, n, idx[uu], idx[vv]);
	return ans;
}
int main(){
	cin>>n>>m;
	for(int i=1; i<=n; i++)	scanf("%d", &w[i]);
	for(int i=1; i<n; i++){
		scanf("%d %d", &uu, &vv);
		add_edge(uu, vv);
		add_edge(vv, uu);
	}
	dfs1(1, 0);
	dfs2(1, 1);
	sgt.build(1, 1, n);
	while(m--){
		scanf("%s", ss);
		if(ss[0]=='C'){
			scanf("%d %d %d", &uu, &vv, &ww);
			changeRange(uu, vv, ww);
		}
		else{
			scanf("%d %d", &uu, &vv);
			printf("%d
", queryRange(uu, vv));
		}
	}
	return 0;
}