[BZOJ 2243] [SDOI 2011] 染色

#include <iostream>
#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>

using namespace std;

const int MaxN = 100000 + 5;

int n, m, Index;
int Father[MaxN], Son[MaxN], Depth[MaxN], Top[MaxN], A[MaxN], Size[MaxN], Pos[MaxN];
int Color[MaxN];

struct Edge
{
	int v;
	Edge *Next;
} E[MaxN * 2], *P = E, *Point[MaxN];

inline void AddEdge(int x, int y) {
	++P; P -> v = y;
	P -> Next = Point[x]; Point[x] = P;
}

int DFS_1(int x, int Dep, int Fa) {
	Depth[x] = Dep; 
	Father[x] = Fa;
	Size[x] = 1;
	int SonSize, MaxSonSize;
	SonSize = MaxSonSize = 0;
	for (Edge *j = Point[x]; j; j = j -> Next) {
		if (j -> v == Fa) continue;
		SonSize = DFS_1(j -> v, Dep + 1, x);
		if (SonSize > MaxSonSize) {
			Son[x] = j -> v;
			MaxSonSize = SonSize;
		}
		Size[x] += SonSize;
	}
	return Size[x];
}

void DFS_2(int x) {
	if (x == 0) return;
	if (x == Son[Father[x]]) Top[x] = Top[Father[x]];
	else Top[x] = x;
	Pos[x] = ++Index;
	Color[Pos[x]] = A[x];
	DFS_2(Son[x]);
	for (Edge *j = Point[x]; j; j = j -> Next) {
		if (j -> v == Father[x] || j -> v == Son[x]) continue;
		DFS_2(j -> v);
	}
}

struct ES 
{
	int Lx, Rx, Tx;
	ES() {}
	ES(int a, int b, int c) {
		Lx = a; Rx = b; Tx = c;
	}
} T[MaxN * 4], D[MaxN * 4];

ES Plus(ES e1, ES e2) {
	ES ret;
	if (e1.Tx == 0) return e2;
	if (e2.Tx == 0) return e1;
	ret = ES(e1.Lx, e2.Rx, e1.Tx + e2.Tx);
	if (e1.Rx == e2.Lx) --ret.Tx;
	return ret;
}

inline void Update(int x) {
	T[x] = Plus(T[x << 1], T[x << 1 | 1]);
}

void Plant_Tree(int x, int s, int t) {
	if (s == t) {
		T[x] = ES(Color[s], Color[s], 1);
		D[x] = ES(0, 0, 0);
		return;
	}
	int m = (s + t) >> 1;
	Plant_Tree(x << 1, s, m);
	Plant_Tree(x << 1 | 1, m + 1, t);
	Update(x);
}

inline void Paint(int x, ES Dlt) {
	T[x] = Dlt;
	D[x] = Dlt;
}

inline void PushDown(int x) {
	if (D[x].Tx == 0) return;
	Paint(x << 1, D[x]);
	Paint(x << 1 | 1, D[x]);
	D[x] = ES(0, 0, 0);
}

void Change_Tree(int x, int s, int t, int l, int r, int Col) {
	if (l <= s && r >= t) {
		ES Temp = ES(Col, Col, 1);
		Paint(x, Temp);
		return;	
	}
	PushDown(x);
	int m = (s + t) >> 1;
	if (l <= m) Change_Tree(x << 1, s, m, l, r, Col);
	if (r >= m + 1) Change_Tree(x << 1 | 1, m + 1, t, l, r, Col);
	Update(x);
}

void Change_Color(int x, int y, int z) {
	int fx, fy;
	while (true) {
		fx = Top[x]; fy = Top[y];
		if (Depth[fx] < Depth[fy]) {
			swap(fx, fy);
			swap(x, y);
		}
		if (fx == fy) {
			if (Pos[x] < Pos[y]) Change_Tree(1, 1, n, Pos[x], Pos[y], z);
			else Change_Tree(1, 1, n, Pos[y], Pos[x], z);
			break;
		}
		else {
			Change_Tree(1, 1, n, Pos[fx], Pos[x], z);
			x = Father[fx];
		}
	}
}

ES Query_Tree(int x, int s, int t, int l, int r) {
	if (l <= s && r >= t) return T[x];
	PushDown(x);
	int m = (s + t) >> 1;
	ES ret = ES(0, 0, 0);
	if (l <= m) ret = Plus(ret, Query_Tree(x << 1, s, m, l, r)); 
	if (r >= m + 1) ret = Plus(ret, Query_Tree(x << 1 | 1, m + 1, t, l, r));
	return ret;
}

int Query(int x, int y) {
	int fx, fy, ret;
	ES Ansx, Ansy, Temp;
	Ansx = Ansy = ES(0, 0, 0);
	while (true) {
		fx = Top[x]; fy = Top[y];
		if (Depth[fx] < Depth[fy]) {
			swap(fx, fy);
			swap(x, y);
			swap(Ansx, Ansy);
		}
		if (fx == fy) {
			ret = Ansx.Tx + Ansy.Tx;
			if (Pos[x] < Pos[y]) {
				Temp = Query_Tree(1, 1, n, Pos[x], Pos[y]);
				ret += Temp.Tx;
				if (Ansy.Lx == Temp.Rx) --ret;
				if (Ansx.Lx == Temp.Lx) --ret;
			}
			else {
				Temp = Query_Tree(1, 1, n, Pos[y], Pos[x]);
				ret += Temp.Tx;
				if (Ansx.Lx == Temp.Rx) --ret;
				if (Ansy.Lx == Temp.Lx) --ret;
			}
			break;
		}
		else {
			Ansx = Plus(Query_Tree(1, 1, n, Pos[fx], Pos[x]), Ansx);
			x = Father[fx];
		}
	}
	return ret;
}

int main() 
{
	scanf("%d%d", &n, &m);
	for (int i = 1; i <= n; ++i) scanf("%d", &A[i]);
	int a, b, c;
	for (int i = 1; i <= n - 1; ++i) {
		scanf("%d%d", &a, &b);
		AddEdge(a, b);
		AddEdge(b, a);
	}
	DFS_1(1, 0, 0);
	Index = 0;
	DFS_2(1);
	Plant_Tree(1, 1, n);
	char ch;
	int Ans;
	for (int i = 1; i <= m; ++i) {
		ch = ‘#‘;
		while (ch != ‘C‘ && ch != ‘Q‘) ch = getchar();
		if (ch == ‘C‘) {
			scanf("%d%d%d", &a, &b, &c);
			Change_Color(a, b, c);
		}
		else {
			scanf("%d%d", &a, &b);
			Ans = Query(a, b);
			printf("%d\n", Ans);
		}
	}
	return 0;
}