D - 染色 HYSBZ - 2243 树链剖分

#include<iostream>
#include<vector>
#include<algorithm>
#include<string>
#include<stack>
#include<cstring>
#include<set>
#include<iterator>
#include<list>
#include<deque>
#include<queue>
#include<map>
#include<cmath>
#include<sstream>
#include<cstdio>
#include<stdio.h>
#include<string.h>
using namespace std;
typedef long long ll;
const int N = 1E5 + 5;
const ll inf = 1e18;
#define lson(x) (x<<1)
#define rson(x) (x<<1|1)
struct Sementtree
{
	int l, r;
	int add, lc, rc;
	ll num;
#define l(x) tree[x].l
#define r(x) tree[x].r
#define add(x) tree[x].add
#define	lc(x) tree[x].lc
#define rc(x) tree[x].rc
#define num(x) tree[x].num
}tree[N << 2];
int n, m;
int head[N], ver[2 * N], Next[2 * N], tot;
int top[N];
int deep[N];
int num[N];
int p[N];
int fp[N];
int fa[N];
int son[N];
int a[N];
int pos = 1;
void pushup(int p)
{
	lc(p) = lc(lson(p));
	rc(p) = rc(rson(p));
	num(p) = num(lson(p)) + num(rson(p)) - (rc(lson(p)) == lc(rson(p)));
}
void build(int p, int l, int r)
{
	l(p) = l, r(p) = r, add(p) = -1;
	if (l == r)
	{
		rc(p) = lc(p) = a[fp[l]];
		num(p) = 1;	
		return;
	}
	int mid = (l + r) / 2;
	build(lson(p), l, mid);
	build(rson(p), mid + 1, r);
	pushup(p);
	//cout << l(p) << " " << r(p) << " "<<num(p) << endl;
}
void spread(int p)
{
	if (l(p) == r(p))
		return;
	if (add(p) != -1)
	{
		num(lson(p)) = num(rson(p)) = 1;
		lc(lson(p)) = rc(lson(p)) = lc(rson(p)) = rc(rson(p)) = add(p);
		add(lson(p)) = add(rson(p)) = add(p);
		add(p) = -1;
	}
}

void change(int p, int l, int r ,int val)
{
	if (l <= l(p) && r >= r(p))
	{
		num(p) = 1, lc(p) = rc(p) = add(p) = val;
		return;
	}
	spread(p);
	int mid = (l(p) + r(p)) / 2;
	if (l <= mid)
		change(lson(p), l, r, val);
	if (r > mid)
		change(rson(p), l, r, val);
	pushup(p);
}
void update(int p, int x, int val)
{
	if (l(p) == x && r(p) == x)
	{
		lc(p) = rc(p) = val;
		add(p) = -1;
		return;
	}
	spread(p);
	int mid = (l(p) + r(p)) / 2;
	if (x <= mid)
	{
		update(lson(p), x, val);
	}
	else
	{
		update(rson(p), x, val);
	}
	pushup(p);
}
ll ask(int p, int l, int r)
{
	//cout << num(p) << endl;
	if (l <= l(p) && r >= r(p))
		return num(p);
	spread(p);
	int mid = (l(p) + r(p)) / 2;
	ll ans = 0, t1 = -1, t2 = -1;
	if (l <= mid)
	{
		t1 = rc(lson(p));
		ans += ask(lson(p), l, r);
	}
	if (r > mid)
	{
		t2 = lc(rson(p));
		ans += ask(rson(p), l, r);
	}
	pushup(p);	
	if (t1 == t2 && t1 != -1)
		ans--;
	return ans;
}
ll ask(int p, int l, int r, int x)
{
	//cout << l << " " << r << " " << lc(p) << " " << rc(p) << endl;
	if (l == r)
	{
		//cout << endl;
		return lc(p);
	}
	int mid = l + r >> 1;
	spread(p);
	if (x <= mid) return ask(lson(p), l, mid, x);
	else return ask(rson(p), mid + 1, r, x);
}
void init()
{
	pos = 1;
	tot = 0;
	memset(head, 0, sizeof(head));
	memset(son, -1, sizeof(son));
}
void addedge(int a, int b)
{
	ver[++tot] = b, Next[tot] = head[a], head[a] = tot;
	ver[++tot] = a, Next[tot] = head[b], head[b] = tot;
}
void dfs1(int u, int pre, int d)
{
	deep[u] = d;
	fa[u] = pre;
	num[u] = 1;
	for (int i = head[u]; i; i = Next[i])
	{
		int v = ver[i];
		if (v != pre)
		{
			dfs1(v, u, d + 1);
			num[u] += num[v];
			if (son[u] == -1 || num[v] > num[son[u]])
			{
				son[u] = v;
			}
		}
	}
}
void getpos(int u, int sp)
{
	top[u] = sp;
	p[u] = pos++;
	fp[p[u]] = u;
	if (son[u] == -1)
		return;
	//cout << u << " " << son[u] << endl;
	getpos(son[u], sp);
	for (int i = head[u]; i; i = Next[i])
	{
		int v = ver[i];
		if (v != son[u] && v != fa[u])
		{
			getpos(v, v);
		}
	}
}

void Query(int u, int v, int flag, int val = 0)
{
	int f1 = top[u], f2 = top[v];
	ll ans = 0;
	while (f1 != f2)
	{
		if (deep[f1] < deep[f2])
		{
			swap(f1, f2);
			swap(u, v);
		}
		if (flag)
		{
			//cout << ask(1, p[f1], p[u])<<" "<< ask(1, p[fa[f1]]) << " " << ask(1, p[f1])<< endl;
			ans += ask(1, p[f1], p[u]) - (ask(1,1,n, p[fa[f1]]) == ask(1,1,n, p[f1]));
		}
		else
			change(1, p[f1], p[u], val);
		u = fa[f1];
		f1 = top[u];
	}
	if (deep[u] > deep[v])
		swap(u, v);
	if (flag)
	{
		ans += ask(1, p[u], p[v]);
		cout << ans << endl;
	}
	else
		change(1, p[u], p[v], val);
}
int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0), cout.tie(0);
	cin >> n >> m;
	for (int i = 1; i <= n; i++)
	{
		cin >> a[i];
	}
	init();
	int x, y;
	for (int i = 0; i < n - 1; i++)
	{	
		cin >> x >> y;
		addedge(x, y);
	}
	dfs1(1, 0, 0);
	getpos(1, 1);
	/*for (int i = 1; i <= n; i++)
	{
		cout << p[i] << " ";
	}
	cout << endl;*/
	build(1, 1, n);
	char c;
	while (m--)
	{
		cin >> c;
		if (c == 'Q')
		{
			cin >> x >> y;
			Query(x, y, 1);
		}
		else
		{
			int t;
			cin >> x >> y >> t;
			Query(x, y, 0, t);
		}
	}
	return 0;
}