BZOJ-2243 [SDOI2011]染色

树链剖分模版题。

#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <iostream>
#include <cctype>
#include <cmath>
#define rep(i, l, r) for(int i=l; i<=r; i++)
#define clr(x, c) memset(x, c, sizeof(x))
#define travel(x) for(edge *p=fir[x]; p; p=p->n)
#define k(x) Key[x]
#define t(x) Tree[x]
#define s(x) Size[x]
#define b(x) Belong[x]
#define low(x) Lower[x]
#define dep(x) Deep[x]
#define h(x) Head[x]
#define l(x) Left[x]
#define r(x) Right[x]
#define w(x) Where[x]
#define maxn 100009
#define maxp 500009
#define inf 0x7fffffff
using namespace std;
inline int read()
{
	int x=0, f=1; char ch=getchar();
	while (!isdigit(ch)) {if (ch=='-') f=-1; ch=getchar();}
	while (isdigit(ch)) x=x*10+ch-'0', ch=getchar();
	return x*f;
}
struct edge{int y; edge *n;} e[maxn*2], *fir[maxn], *pt=e;
inline void AddE(int x, int y)
{
	pt->y=y, pt->n=fir[x], fir[x]=pt++;
	pt->y=x, pt->n=fir[y], fir[y]=pt++;
}
int Tree[maxn], Size[maxn], Belong[maxn], Lower[maxn], Deep[maxn], Head[maxn], Where[maxn], cnt, Key[maxn];
int Left[maxp], Right[maxp], sum[maxp], lc[maxp], rc[maxp], z, L, R, C;
int n, m, d[maxn], h[maxn], now, nowrc, nowlc;
bool tag[maxp];


void dfs(int x)
{
	int maxs=-inf, maxl=0;
	travel(x) if (p->y!=h[x])
	{
		h[p->y]=x, d[p->y]=d[x]+1; dfs(p->y);
		if (s(b(p->y))>maxs) maxs=s(b(p->y)), maxl=p->y;
	}
	if (maxl)
		b(x)=b(maxl), w(x)=w(maxl)+1, s(b(x))++, dep(b(x))--;
	else
		cnt++, b(x)=cnt, w(x)=1, s(cnt)=1, low(cnt)=x, dep(cnt)=d[x];
	travel(x) if (p->y!=h[x] && p->y!=maxl) h(b(p->y))=x;
}
void BuildT(int&k, int l, int r, int t)
{
	if (l==r){sum[t]=1, lc[t]=rc[t]=k(k), k=h[k]; return;}
	int mid=(l+r)>>1;
	BuildT(k, l, mid, l(t)=++z), BuildT(k, mid+1, r, r(t)=++z);
	sum[t]=sum[l(t)]+sum[r(t)], lc[t]=lc[l(t)], rc[t]=rc[r(t)];
	if (rc[l(t)]==lc[r(t)]) sum[t]--;
}
inline void Build()
{
	dfs(1); h(b(1))=0;
	rep(i, 1, cnt) now=low(i), BuildT(now, 1, s(i), t(i)=++z);
}


inline void pushdown(int t)
{
	tag[t]=0, tag[l(t)]=tag[r(t)]=1;
	lc[l(t)]=rc[l(t)]=lc[r(t)]=rc[r(t)]=lc[t];
	sum[l(t)]=sum[r(t)]=1;
}
void Edit(int l, int r, int t)
{
	if (L<=l && r<=R){tag[t]=true, sum[t]=1, lc[t]=rc[t]=C; return;}
	int mid=(l+r)>>1;
	if (tag[t]) pushdown(t);
	if (L<=mid) Edit(l, mid, l(t));
	if (mid<R) Edit(mid+1, r, r(t));
	sum[t]=sum[l(t)]+sum[r(t)], lc[t]=lc[l(t)], rc[t]=rc[r(t)];
	if (rc[l(t)]==lc[r(t)]) sum[t]--;
}
inline void Change(int x, int y)
{
	while (b(x)!=b(y))
	{
		if (dep(b(x))<dep(b(y))) swap(x, y);
		L=w(x), R=s(b(x)), Edit(1, s(b(x)), t(b(x))), x=h(b(x));
	}
	if (d[x]<d[y]) swap(x, y);
	L=w(x), R=w(y), Edit(1, s(b(x)), t(b(x)));
}


void Query(int l, int r, int t)
{
	if (L<=l && r<=R) {now+=sum[t]; if (nowrc==lc[t]) now--; nowrc=rc[t]; return;}
	if (tag[t]) pushdown(t);
	int mid=(l+r)>>1;
	if (L<=mid) Query(l, mid, l(t));
	if (mid<R) Query(mid+1, r, r(t));
}
inline void Qsum(int x, int y)
{
	now=0; nowrc=nowlc=-1; while (b(x)!=b(y))
	{
		if (dep(b(x))<dep(b(y))) swap(x, y), swap(nowrc, nowlc);
		L=w(x), R=s(b(x)), Query(1, s(b(x)), t(b(x))), x=h(b(x));
	}
	if (d[x]<d[y]) swap(x, y), swap(nowrc, nowlc);
	L=w(x), R=w(y), Query(1, s(b(x)), t(b(x))); 
	if (nowrc==nowlc) now--;
	printf("%d
", now);
}


int main()
{
	n=read(), m=read();
	rep(i, 1, n) k(i)=read();
	rep(i, 1, n-1) AddE(read(), read());
	Build();
	rep(i, 1, m)
	{
		char ch[5]; scanf("%s", ch); int x=read(), y=read();
		if (ch[0]=='C') C=read(), Change(x, y); else Qsum(x, y);
	}
	return 0;
}