BZOJ2243 [SDOI2011]染色（树链剖分+线段树合并）

题目链接 BZOJ2243

树链剖分 $+$ 线段树

线段树每个节点维护$lc$, $rc$, $s$

$lc$代表该区间的最左端的颜色，$rc$代表该区间的最右端的颜色

$s$代表该区间的所有连续颜色段数（仅考虑该区间）

$lazy$表示延迟信息。

#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b)	for (int i(a); i <= (b); ++i)
#define dec(i, a, b)	for (int i(a); i >= (b); --i)
#define lson		i << 1, L, mid
#define rson		i << 1 | 1, mid + 1, R

typedef long long LL;

const int N = 1e5 + 10;
const int A = 21;

vector <int> v[N];
int n, m;
int c[N];
int f[N][A];
int lazy[N << 2], s[N << 2], lc[N << 2], rc[N << 2];
int h[N], deep[N], sz[N];
int son[N];
int tot = 0;
int top[N];

int LCA(int a, int b){
	if (deep[a] < deep[b]) swap(a, b);
	for (int i = 0,  delta = deep[a] - deep[b]; delta; delta >>= 1, ++i) if (delta & 1) a = f[a][i];
	if (a == b) return a;
	dec(i, 19, 0) if (f[a][i] != f[b][i]) a = f[a][i], b = f[b][i];
	return f[a][0];
}

inline void pushup(int i){
	lc[i] = lc[i << 1];
	rc[i] = rc[i << 1 | 1];
	if (rc[i << 1] ^ lc[i << 1 | 1]) s[i] = s[i << 1] + s[i << 1 | 1];
	else s[i] = s[i << 1] + s[i << 1 | 1] - 1;
}

inline void pushdown(int i, int L, int R){
	int tmp = lazy[i];
	if (tmp == -1 || L == R) return;

	s[i << 1]      = s[i << 1 | 1]    = 1;
	lazy[i << 1]   = lazy[i << 1 | 1] = tmp;

	lc[i << 1]     = rc[i << 1]       = tmp;
	lc[i << 1 | 1] = rc[i << 1 | 1]   = tmp;
	lazy[i] = -1;
}	

void dfs1(int x, int fa, int dep){
	sz[x] = 1;
	deep[x] = dep;

	if (fa){
                f[x][0] = fa;
		for (int i = 0; f[f[x][i]][i]; ++i) f[x][i + 1] = f[f[x][i]][i];
	}

	int ct = (int)v[x].size();

	rep(i, 0, ct - 1){
		int u = v[x][i];
		if (u == fa) continue;
		dfs1(u, x, dep + 1);
		sz[x] += sz[u];
		if (sz[son[x]] < sz[u]) son[x] = u;
	}
}

void dfs2(int x, int tp){
	h[x] = ++tot;
	top[x] = tp;
	if (son[x]) dfs2(son[x], tp);

	int ct = (int)v[x].size();
	rep(i, 0, ct - 1){
		int u = v[x][i];
		if (u == f[x][0] || u == son[x]) continue;
		dfs2(u, u);
	}
}

void build(int i, int L, int R){
	s[i] = 1;
	lazy[i] = -1;

	if (L == R) return;
	int mid = (L + R) >> 1;

	build(lson);
	build(rson);
}

void change(int i, int L, int R, int l, int r, int val){
	pushdown(i, L, R);
	if (L == l && R == r){
		lc[i] = rc[i] = val;
		s[i]  = 1;
		lazy[i] = val;
		return;
	}

	int mid = (L + R) >> 1;
	if (r <= mid) change(lson, l, r, val);
	else if (l > mid) change(rson, l, r, val);
	else{
		change(lson, l, mid, val);
		change(rson, mid + 1, r, val);
	}

	pushup(i);
}

int query(int i, int L, int R, int l, int r){
	pushdown(i, L, R);
	if (L == l && R == r) return s[i];

	int mid = (L + R) >> 1;
	if (r <= mid) return query(lson, l, r);
	else if (l > mid) return query(rson, l, r);
	else{
		int tmp = 1;
		if (rc[i << 1] ^ lc[i << 1 | 1]) tmp = 0;
		return query(lson, l, mid) + query(rson, mid + 1, r) - tmp;
	}
}

int getcolor(int i, int L, int R, int x){
	pushdown(i, L, R);
	if (L == R) return lc[i];
	int mid = (L + R) >> 1;
	if (x <= mid) return getcolor(lson, x);
	else return getcolor(rson, x);
}

int solvesum(int x, int tp){
	int ret = 0;
	for (; top[x] ^ top[tp] ;){
		ret += query(1, 1, n, h[top[x]], h[x]);
		if (getcolor(1, 1, n, h[top[x]]) == getcolor(1, 1, n, h[f[top[x]][0]])) --ret;
		x = f[top[x]][0];
	}

	ret += query(1, 1, n, h[tp], h[x]);
	return ret;
}

void solvechange(int x, int tp, int val){
	for (; top[x] ^ top[tp]; ){
		change(1, 1, n, h[top[x]], h[x], val);
		x = f[top[x]][0];
	}

	change(1, 1, n, h[tp], h[x], val);
}

void solve(){
	int x, y, z;
	dfs1(1, 0, 0);
	dfs2(1, 1);
	build(1, 1, n);
	rep(i, 1, n) change(1, 1, n, h[i], h[i], c[i]);

	rep(i, 1, m){
		char ch[10];
		scanf("%s", ch);
		if (ch[0] == 'Q'){
			scanf("%d%d", &x, &y);
			int t = LCA(x, y);
			printf("%d\n", solvesum(x, t) + solvesum(y, t) - 1);
		}
		else{
			scanf("%d%d%d", &x, &y, &z);
			int t = LCA(x, y);
			solvechange(x, t, z);
			solvechange(y, t, z);
		}
	}
}

void init(){
	scanf("%d%d", &n, &m);
	rep(i, 1, n) scanf("%d", c + i);
	rep(i, 2, n){
		int x, y;
		scanf("%d%d", &x, &y);
		v[x].push_back(y);
		v[y].push_back(x);
	}
}

int main(){
	init();
	solve();
	return 0;
}