洛谷P2486 染色(树剖入门)

题意：

给定一棵有$n$个节点的无根树和$m$个操作，操作有$2$类：
$1、$将节点$a$到节点$b$路径上所有点都染成颜色$c$；
$2、$询问节点$a$到节点$b$路径上的颜色段数量（连续相同颜色被认为是同一段）

分析：

树剖之后用线段树维护颜色就好了。
维护的操作：线段树维护每个区间的左颜色$lcol$和右颜色$lcol$，和区间的颜色数$w$，$lazy$就是记录这个区间的颜色。
线段树每次向上更新时判断节点的左儿子的右颜色和右二子的左颜色是否相同，如果是则区间颜色数减$1$，不是的话正常更新即可。
询问的时候也要注意一点：每次要判断链顶端点和链顶端点父亲的颜色是否相同，如果是答案就要减$1$。

#include <bits/stdc++.h>
using namespace std;

typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int, int> pii;

const int maxn  = 1e5 + 5;
const int maxm  = 100 + 5;
const int inf   = 0x3f3f3f3f;
const LL  mod   = 1e9 + 7;//19260817
const double pi = acos(-1.0); 

int n, m, x, y, z, cnt, tot, w[maxn], head[maxn];
int id[maxn], pre[maxn], son[maxn], size[maxn], deep[maxn], top[maxn], nw[maxn];
char str[20];

struct node{
	int to, next;
}edge[maxn << 1];

struct qwq{
	int w, lcol, rcol, lazy;
}a[maxn << 2];

void pushup(int k){
	a[k].w = a[k << 1].w + a[k << 1 | 1].w;
	a[k].lcol = a[k << 1].lcol;
	a[k].rcol = a[k << 1 | 1].rcol;
	if(a[k << 1].rcol == a[k << 1 | 1].lcol) a[k].w--;
}

void pushdown(int k){
	if(a[k].lazy){
		a[k << 1].lazy = a[k << 1 | 1].lazy = a[k].lazy;
		a[k << 1].w = a[k << 1 | 1].w = 1;
		a[k << 1].lcol = a[k << 1].rcol = a[k].lazy;
		a[k << 1 | 1].lcol = a[k << 1 | 1].rcol = a[k].lazy;
		a[k].lazy = 0;
	}
}

void build(int k, int l, int r){
	if(l == r){
	//	printf("build: l = %d nw[l] = %d\n", l, nw[l]);
		a[k].lcol = a[k].rcol = nw[l];
		a[k].w = 1;
		return ;
	}
	int mid = (l + r) >> 1;
	build(k << 1, l, mid);
	build(k << 1 | 1, mid + 1, r);
	pushup(k);
}

void update(int k, int l, int r, int x, int y, int z){
	if(l >= x && r <= y){
		a[k].w = 1;
		a[k].lcol = a[k].rcol = z;
		a[k].lazy = z;
		return ;
	}
	pushdown(k);
	int mid = (l + r) >> 1;
	if(x <= mid) update(k << 1, l, mid, x, y, z);
	if(y > mid) update(k << 1 | 1, mid + 1, r, x, y, z);
	pushup(k);
}

int query(int k, int l, int r, int x, int y){
	if(l >= x && r <= y){
		return a[k].w;
	}
	pushdown(k);
	int mid = (l + r) >> 1, ans = 0;
	if(x <= mid) ans += query(k << 1, l, mid, x, y);
	if(y > mid) ans += query(k << 1 | 1, mid + 1, r, x, y);
	if(x <= mid && y > mid && a[k << 1].rcol == a[k << 1 | 1].lcol) ans--;
	return ans;
}

int queryColor(int k, int l, int r, int x){
	if(l == x && r == x){
		return a[k].lcol;
	}
	pushdown(k);
	int mid = (l + r) >> 1;
	if(x <= mid) return queryColor(k << 1, l, mid, x);
	if(x > mid) return queryColor(k << 1 | 1, mid + 1, r, x);
}

void addedge(int u, int v){
	edge[++cnt].to = v;
	edge[cnt].next = head[u];
	head[u] = cnt;
}

void dfs1(int x, int f, int d){
	pre[x] = f;
	deep[x] = d;
	size[x] = 1;
	int MAX = -1;
	for(int i = head[x]; i; i = edge[i].next){
		int v = edge[i].to;
		if(v == f) continue;
		dfs1(v, x, d + 1);
		size[x] += size[v];
		if(size[v] > MAX) MAX = size[v], son[x] = v;
	}
}

void dfs2(int x, int topf){
	id[x] = ++tot;
	nw[tot] = w[x];
	top[x] = topf;
	if(!son[x]) return ;
	dfs2(son[x], topf);
	for(int i = head[x]; i; i = edge[i].next){
		int v = edge[i].to;
		if(v == pre[x] || v == son[x]) continue;
		dfs2(v, v);
	}
} 

int queryCol(int x, int y){
	int ans = 0;
	while(top[x] != top[y]){
		if(deep[top[x]] < deep[top[y]]) swap(x, y);
		ans += query(1, 1, n, id[top[x]], id[x]);
		if(queryColor(1, 1, n, id[top[x]]) == queryColor(1, 1, n, id[pre[top[x]]])) ans--;
		x = pre[top[x]];
	}
	if(deep[x] > deep[y]) swap(x, y);
	ans += query(1, 1, n, id[x], id[y]);
	return ans;
}

void updateCol(int x, int y, int z){
	while(top[x] != top[y]){
		if(deep[top[x]] < deep[top[y]]) swap(x, y);
		update(1, 1, n, id[top[x]], id[x], z);
		x = pre[top[x]];
	}
	if(deep[x] > deep[y]) swap(x, y);
	update(1, 1, n, id[x], id[y], z);
}

int main(){
	scanf("%d %d", &n, &m);
	for(int i = 1; i <= n; i++) scanf("%d", w + i);
	for(int i = 1; i <= n - 1; i++){
		scanf("%d %d", &x, &y);
		addedge(x, y), addedge(y, x);
	}
	dfs1(1, -1, 1);
	dfs2(1, 1);
	build(1, 1, n);
	while(m--){
		scanf("%s", str);
		if(str[0] == 'Q'){
			scanf("%d %d", &x, &y);
			printf("%d\n", queryCol(x, y));
		}else if(str[0] == 'C'){
			scanf("%d %d %d", &x, &y, &z);
			updateCol(x, y, z);
		}
	}
    return 0;
}