BZOJ 2243: [SDOI2011]染色 树链剖分

Description

给定一棵有n个节点的无根树和m个操作，操作有2类：
1、将节点a到节点b路径上所有点都染成颜色c；
2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），
如“112221”由3段组成：“11”、“222”和“1”。
请你写一个程序依次完成这m个操作。

Input

第一行包含2个整数n和m，分别表示节点数和操作数；
第二行包含n个正整数表示n个节点的初始颜色
下面 行每行包含两个整数x和y，表示x和y之间有一条无向边。
下面 行每行描述一个操作：
“C a b c”表示这是一个染色操作，把节点a到节点b路径上所有点（包括a和b）都染成颜色c；
“Q a b”表示这是一个询问操作，询问节点a到节点b（包括a和b）路径上的颜色段数量。

Output

对于每个询问操作，输出一行答案。

Sample Input

6 5

2 2 1 2 1 1

1 2

1 3

2 4

2 5

2 6

Q 3 5

C 2 1 1

Q 3 5

C 5 1 2

Q 3 5

Sample Output

3

1

2
HINT
数N<=10^{5，操作数M<=10}5，所有的颜色C为整数且在[0, 10^9]之间。

先建一个树链剖分，然后写一个线段树就好了。注意合并的时候颜色的变化。

#include<bits/stdc++.h>
#define mem(x,v) memset(x,v,sizeof(x))
#define lson now << 1
#define rson now << 1 | 1
using namespace std;
const int N = 2e5+100;
int n,m;
struct node
{
	int lazy,val;
	int lc,rc;	
}tree[N*4];
int cnt,Head[N],Next[N*2],to[N*2],w[N],wt[N];
int son[N],id[N],tim,size[N],top[N],dep[N],f[N];
int ft[N][20];
void Add_edge(int u, int v){
	to[++cnt] = v;
	Next[cnt] = Head[u];
	Head[u] = cnt;
}
void dfs1(int u, int fa, int deep){
	dep[u] = deep; f[u] = fa; size[u] = 1;
	ft[u][0] = fa;
	for (int i = 1; i < 20; i++)
		ft[u][i] = ft[ft[u][i-1]][i-1];
	for (int i = Head[u]; i; i = Next[i]){
		int v = to[i];
		if (v == fa) continue;
		dfs1(v, u, deep + 1);
		size[u] += size[v];
		if (size[v] > size[son[u]]) son[u] = v;
	}
}
void dfs2(int u, int topf){
	id[u] = ++tim;
	wt[tim] = w[u];
	top[u] = topf;
	if (!son[u]) return;
	dfs2(son[u],topf);
	for (int i = Head[u]; i; i = Next[i]){
		int v = to[i];
		if (v == f[u] || v == son[u]) continue;
		dfs2(v,v);
	}
}
void pushdown(int now){
	if (tree[now].lazy == -1) return;
	tree[lson].lc = tree[lson].rc = tree[now].lazy;
	tree[rson].lc = tree[rson].rc = tree[now].lazy;
	tree[lson].lazy = tree[rson].lazy = tree[now].lazy;
	tree[lson].val = tree[rson].val =1;
	tree[now].lazy = -1;
}
void Build(int now, int l, int r){
	tree[now].lazy = -1;   //lazy 标记，一开始算是没有颜色。
	if (l + 1 == r){
		tree[now].lc = wt[l];
		tree[now].rc = wt[l];
		tree[now].val = 1;
		return;
	}
	int mid = (l + r) >> 1;
	Build(lson,l,mid); Build(rson,mid,r);
	tree[now].val = tree[lson].val + tree[rson].val;
	if (tree[lson].rc == tree[rson].lc) tree[now].val--;  //合并的时候判断，中间的点颜色是不是一样的。如果是，颜色数要 -1
	tree[now].lc = tree[lson].lc;
	tree[now].rc = tree[rson].rc;
}

void Insert(int now, int l,int r, int a, int b, int k){
	if (a <= l && b >= r-1){
		tree[now].val = 1;
		tree[now].lc = tree[now].rc = k;
		tree[now].lazy = k;
		return;
	}
	int mid = (l + r) >> 1;
	pushdown(now);
	if (a < mid)  Insert(lson,l,mid,a,b,k);
	if (b >= mid) Insert(rson,mid,r,a,b,k);
	tree[now].val = tree[lson].val + tree[rson].val;
	if (tree[lson].rc == tree[rson].lc) tree[now].val--;
	tree[now].lc = tree[lson].lc;
	tree[now].rc = tree[rson].rc;
}
int Query(int now,int l, int r, int a, int b){
	int ans = 0;
	if (a <= l && b >= r - 1){
		ans += tree[now].val;
		return ans;
	}
	pushdown(now);
	int mid = (l + r) >> 1;
	if (a < mid) ans += Query(lson,l,mid,a,b);
	if (b >= mid) ans += Query(rson,mid,r,a,b);
	if ((a < mid) && (b >= mid)){
		if (tree[lson].rc == tree[rson].lc) ans--;
	}
	return ans;
}
int lca(int x, int y){
	if (dep[x] < dep[y]) swap(x,y);
	for (int i = 19; i >= 0; i--)
		if (dep[ft[x][i]] >= dep[y]) x = ft[x][i];
	if (x == y) return x;
	for (int i = 19; i >= 0; i--)
		if (ft[x][i] != ft[y][i]){
			x = ft[x][i]; y = ft[y][i];
		}
	return ft[x][0];
}
void upRange(int x, int y, int k){
	while(top[x] != top[y]){
		Insert(1,1,n+1,id[top[x]],id[x],k);
		x = f[top[x]];
	}
	Insert(1,1,n+1,id[y],id[x],k);
}
int ask_col(int now,int l, int r, int k){
	if (l + 1 == r){
		return tree[now].lc;
	}
	int mid = (l + r) / 2;
	pushdown(now);
	if (k < mid) return ask_col(lson,l,mid,k); else
	return ask_col(rson,mid,r,k);
}
int ask_Range(int x, int y){
	int ans = 0;
	//这个地方知道了两个点的最近公共祖先是其中 y 点，所以就不需要比较两个点的深度了。
	while(top[x] != top[y]){
		ans += Query(1,1,n+1,id[top[x]],id[x]);
		if (ask_col(1,1,1+n,id[top[x]]) == ask_col(1,1,1+n,id[f[top[x]]])) ans--;  
		//两个相连接的路径上，判断相邻的两个点的颜色是不是一样的，。 
		x = f[top[x]];
	}
	ans += Query(1,1,1+n,id[y],id[x]);
	return ans;
}

int main(){
	scanf("%d%d",&n,&m);
	for (int i = 1; i <= n; i++) 
		scanf("%d",&w[i]);
	int x,y,z;
	for (int i = 1; i < n; i++){
		scanf("%d%d",&x,&y);
		Add_edge(x,y); Add_edge(y,x);
	}
	dfs1(1,0,1); 
	dfs2(1,1);
	Build(1,1,n+1);
	char ch[10];
	int ans;
	for (int i = 0; i < m; i++){
		scanf("%s",ch);  //读入错了，改了好长时间。 
		if (ch[0] == 'Q'){
			scanf("%d%d",&x,&y);
			z = lca(x,y); //两个点不一定在一条路上，找最高的点。
			ans = ask_Range(x,z); // 分成两条路。
			ans += ask_Range(y,z);
			printf("%d\n",ans-1); //最近公共祖先一定是重复计算，所以 -1。
		} else {
			scanf("%d%d%d",&x,&y,&z);
			int mark = lca(x,y);
			upRange(x,mark,z);
			upRange(y,mark,z);
		}
	}
	return 0;
}


/*

6 5
2 2 1 2 1 1
1 2
1 3
2 4
2 5
2 6
Q 3 5
C 2 1 1
Q 3 5
C 5 1 2
Q 3 5

*/