BZOJ 2243: [SDOI2011]染色 【树剖+线段树维护颜色段数】

题目传送门

线段树存区间最左边和最右边的颜色，合并左右儿子的段数的时候看是否需要减1.
写一个结构体重载+号似乎很方便，统计树链的时候直接加就可以了。
询问的时候两条链要分别跑，不能用swap.

Code：

#include<cstdio>
#include<cctype>
#include<algorithm>
#define maxn 100005
using namespace std;
char cb[1<<15],*cs,*ct;
#define getc() (cs==ct&&(ct=(cs=cb)+fread(cb,1,1<<15,stdin),cs==ct)?0:*cs++)
inline void read(int &a){
	char c; bool f=0; while(!isdigit(c=getc())) if(c=='-') f=1;
	for(a=c-'0';isdigit(c=getc());a=a*10+c-'0'); if(f) a=-a;
}
const int inf = 1<<30;
int n,m,a[maxn],seq[maxn];
int dfn[maxn],tim,fa[maxn],son[maxn],siz[maxn],top[maxn],dep[maxn];
int fir[maxn],nxt[maxn<<1],to[maxn<<1],tot;
inline void line(int x,int y){nxt[++tot]=fir[x],fir[x]=tot,to[tot]=y;}

void dfs1(int u){
	siz[u]=1;
	for(int i=fir[u],v;i;i=nxt[i]) if((v=to[i])!=fa[u]){
		fa[v]=u,dep[v]=dep[u]+1;
		dfs1(v);siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}
void dfs2(int u,int tp){
	top[u]=tp,seq[dfn[u]=++tim]=a[u];
	if(son[u]) dfs2(son[u],tp);
	for(int i=fir[u];i;i=nxt[i]) if(!dfn[to[i]]) dfs2(to[i],to[i]);
}

#define lc (i<<1)
#define rc (i<<1|1)
struct data{
	int col,cor,s;
	data(){col=cor=-1,s=0;}
	data(int col,int cor,int s):col(col),cor(cor),s(s){}
	data operator + (const data &b)const{
		return data(col,b.cor,s+b.s-(cor==b.col));
	}
};
struct node{
	int tag;
	data v;
	node(){tag=-1;}
	void init(int x){v=data(x,x,1);tag=x;}
}t[maxn<<2];
void upd(int i){
	t[i].v=t[lc].v+t[rc].v;
}
void pushdown(int i){
	if(t[i].tag!=-1){
		t[lc].init(t[i].tag),t[rc].init(t[i].tag);
		t[i].tag=-1;
	}
}
void build(int i,int l,int r){
	if(l==r) {t[i].init(seq[l]);return;}
	int mid=(l+r)>>1;
	build(lc,l,mid);
	build(rc,mid+1,r);
	upd(i);
}
void insert(int i,int l,int r,int x,int y,int c){
	if(x<=l&&r<=y) {t[i].init(c);return;}
	pushdown(i);
	int mid=(l+r)>>1;
	if(x<=mid) insert(lc,l,mid,x,y,c);
	if(y>mid) insert(rc,mid+1,r,x,y,c);
	upd(i);	
}
data query(int i,int l,int r,int x,int y){
	if(x<=l&&r<=y) return t[i].v;
	pushdown(i);
	int mid=(l+r)>>1;
	if(y<=mid) return query(lc,l,mid,x,y);
	if(x>mid) return query(rc,mid+1,r,x,y);
	return query(lc,l,mid,x,y)+query(rc,mid+1,r,x,y);
}

void modify(int u,int v,int c){
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]) swap(u,v);
		insert(1,1,n,dfn[top[u]],dfn[u],c);
		u=fa[top[u]];
	}
	if(dep[u]<dep[v]) swap(u,v);
	insert(1,1,n,dfn[v],dfn[u],c);
}
int ask(int u,int v){
	data retu,retv;
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]) retv=query(1,1,n,dfn[top[v]],dfn[v])+retv,v=fa[top[v]];
		else retu=query(1,1,n,dfn[top[u]],dfn[u])+retu,u=fa[top[u]];
	}
	if(dep[u]<dep[v]) retv=query(1,1,n,dfn[u],dfn[v])+retv;
	else retu=query(1,1,n,dfn[v],dfn[u])+retu;//一开始u，v写反了
	swap(retu.col,retu.cor);
	return (retu+retv).s;//lca在两条链中都包含，所以直接加就可以了
}
int main()
{
	//freopen("H.in","r",stdin);
	int x,y,z;char op;
	read(n),read(m);
	for(int i=1;i<=n;i++) read(a[i]);
	for(int i=1;i<n;i++) read(x),read(y),line(x,y),line(y,x);
	dfs1(1),dfs2(1,1);
	build(1,1,n);
	while(m--){
		while(!isalpha(op=getc()));read(x),read(y);
		if(op=='C') read(z),modify(x,y,z);
		else printf("%d\n",ask(x,y));
	}
}