BZOJ 2243: [SDOI2011]染色 (树剖+线段树)

树链剖分后两个区间合并的时候就判一下相交颜色是否相同来算颜色段数就行了.

CODE

#include <vector>
#include <queue>
#include <cstdio>
#include <cctype>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long LL;
#define ls (i<<1)
#define rs (i<<1|1)
inline void read(int &num) {
	char ch; int flg=1; while(!isdigit(ch=getchar()))if(ch=='-')flg=-flg;
	for(num=0; isdigit(ch); num=num*10+ch-'0',ch=getchar()); num*=flg;
}
const int MAXN = 100005;

int n, q, w[MAXN], dfn[MAXN], seq[MAXN], tmr;
int fir[MAXN], cnt;
struct edge { int to, nxt; }e[MAXN<<1];

inline void add(int u, int v) {
	e[cnt] = (edge){ v, fir[u] }, fir[u] = cnt++;
	e[cnt] = (edge){ u, fir[v] }, fir[v] = cnt++;
}
int dep[MAXN], fa[MAXN], sz[MAXN], top[MAXN], son[MAXN];
void dfs(int u, int ff) {
	dep[u] = dep[fa[u]=ff] + (sz[u]=1);
	for(int v, i = fir[u]; ~i; i = e[i].nxt) 
		if((v=e[i].to) != fa[u]) {
			dfs(v, u), sz[u] += sz[v];
			if(sz[v] > sz[son[u]]) son[u] = v;
		}
}
void dfs2(int u, int tp) {
	top[u] = tp; seq[dfn[u] = ++tmr] = u;
	if(son[u]) dfs2(son[u], tp);
	for(int v, i = fir[u]; ~i; i = e[i].nxt)
		if((v=e[i].to) != son[u] && v != fa[u]) dfs2(v, v);
}
struct node {
	int lc, rc, sum;
	node(){ sum = -1; }
	node(int l, int r, int s):lc(l), rc(r), sum(s){}
	inline node operator +(const node &o)const {
		if(sum == -1) return o;
		if(o.sum == -1) return *this;
		if(rc == o.lc) return node(lc, o.rc, sum+o.sum-1);
		else return node(lc, o.rc, sum+o.sum);
	}
	inline friend node rev(const node &o) { return node(o.rc, o.lc, o.sum); }
}col[MAXN<<2];
int tag[MAXN<<2];
inline void upd(int i) {
	col[i] = col[ls] + col[rs];
}
inline void mt(int i) {
	if(~tag[i]) {
		col[ls] = col[rs] = node(tag[i], tag[i], 1);
		tag[ls] = tag[rs] = tag[i]; //忘了写这个WA到自闭
		tag[i] = -1;
	}
}
void build(int i, int l, int r) {
	tag[i] = -1;
	if(l == r) {
		col[i] = node(w[seq[l]], w[seq[l]], 1);
		return;
	}
	int mid = (l + r) >> 1;
	build(ls, l, mid);
	build(rs, mid+1, r);
	upd(i);
}
void cover(int i, int l, int r, int x, int y, int val) {
	if(l == x && r == y) {
		col[i] = node(val, val, 1);
		tag[i] = val;
		return;
	}
	int mid = (l + r) >> 1;
	mt(i);
	if(y <= mid) cover(ls, l, mid, x, y, val);
	else if(x > mid) cover(rs, mid+1, r, x, y, val);
	else cover(ls, l, mid, x, mid, val), cover(rs, mid+1, r, mid+1, y, val);
	upd(i);
}
node query(int i, int l, int r, int x, int y) {
	if(l == x && r == y) return col[i];
	int mid = (l + r) >> 1; node res;
	mt(i);
	if(y <= mid) res = query(ls, l, mid, x, y);
	else if(x > mid) res = query(rs, mid+1, r, x, y);
	else res = query(ls, l, mid, x, mid) + query(rs, mid+1, r, mid+1, y);
	upd(i);
	return res;
}
inline node Query(int x, int y) {
	node resx, resy;
	while(top[x] != top[y]) {
		if(dep[top[x]] > dep[top[y]]) resx = resx + rev(query(1, 1, n, dfn[top[x]], dfn[x])), x = fa[top[x]];
		else resy = query(1, 1, n, dfn[top[y]], dfn[y]) + resy, y = fa[top[y]];
	}
	if(dfn[x] < dfn[y]) return resx + query(1, 1, n, dfn[x], dfn[y]) + resy;
	else return resx + rev(query(1, 1, n, dfn[y], dfn[x])) + resy;
}
inline void Cover(int x, int y, int val) {
	while(top[x] != top[y]) {
		if(dep[top[x]] < dep[top[y]]) swap(x, y);
		cover(1, 1, n, dfn[top[x]], dfn[x], val);
		x = fa[top[x]];
	}
	if(dfn[x] < dfn[y]) swap(x, y);
	cover(1, 1, n, dfn[y], dfn[x], val);
}
int main() {
	read(n); read(q);
	for(int i = 1; i <= n; ++i) fir[i] = -1, read(w[i]);
	for(int i = 1, x, y; i < n; ++i)
		read(x), read(y), add(x, y);
	dfs(1, 0); dfs2(1, 1);
	build(1, 1, n);
	char s[2]; int x, y, z;
	while(q--) {
		while(!isalpha(s[0]=getchar()));
		read(x), read(y);
		if(s[0] == 'Q') printf("%d\n", Query(x, y).sum);
		else read(z), Cover(x, y, z);
	}
}