[SDOI2011]染色（树链剖分，线段树）

发表于 更新于 分类于 ， ， ， ，
SCOI2011

题目

Description

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

Input

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

Output

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

Sample Input

1
2
3
4
5
6
7
8
9
10
11
12
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

1
2
3
3
1
2

HINT

$N \leq 10^5, M \leq 10^5$，所有的颜色C为整数且在$[0, 10^9]$之间。

解题思路

这道题思路应该是很好想到的——树链剖分+线段树
显然，合并两个颜色段时，如果合并点两侧颜色相同，那么颜色段数量为左右两段数量之和再减一，否则就是它们的和。根据这个基本性质——
线段树维护3个值：颜色段数量(cnt)、左端颜色(lcol)、右端颜色(rcol)，于是tr[id].cnt = tr[lid].cnt + tr[rid].cnt - (tr[lid].rcol == tr[rid].lcol);
树链剖分询问时（即往上“跳”时），记录一下上一次询问的端点颜色，如果本次询问的相应端点颜色在合并点处与上次询问端点颜色相等，就要减一。
至此本题解决。

复杂度$O(n\log ^2n)$

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
# include<cstdio>
# include<algorithm>

using namespace std;

const int N = 100005;
int n, m, a[N], u, v, qa, qb, qc;
char opt[2];

struct Edge{
	int nxt, to;
}edge[N<<1];
int head[N], edgeNum;
void addEdge(int from, int to){
	edge[++edgeNum].nxt = head[from];
	edge[edgeNum].to = to;
	head[from] = edgeNum;
}

int fa[N], dep[N], size[N], son[N];
void dfs1(int x, int f, int depth){
	fa[x] = f, dep[x] = depth, size[x] = 1, son[x] = 0;
	for(int i = head[x]; i; i = edge[i].nxt){
		if(edge[i].to == f)	continue;
		dfs1(edge[i].to, x, depth+1);
		size[x] += size[edge[i].to];
		if(size[edge[i].to] > size[son[x]])	son[x] = edge[i].to;
	}
}
int belong[N], st[N], ed[N], dfsClock, fun[N];
void dfs2(int x, int top){
	st[x] = ++dfsClock, fun[dfsClock] = x;
	belong[x] = top;
	if(son[x])	dfs2(son[x], top);
	for(int i = head[x]; i; i = edge[i].nxt){
		if(edge[i].to == fa[x] || edge[i].to == son[x])	continue;
		dfs2(edge[i].to, edge[i].to);
	}
	ed[x] = dfsClock;
}

# define lid id<<1
# define rid id<<1|1
# define mid ((tr[id].l+tr[id].r)>>1)

struct segTree{
	int l, r, lcol, rcol, cnt, cov;
	segTree(){
		l = r = lcol = rcol = cnt = cov = 0;
	}
}tr[N<<2];

struct {
	inline void pushup(int id){
		tr[id].cnt = tr[lid].cnt + tr[rid].cnt - (tr[lid].rcol == tr[rid].lcol);
		tr[id].lcol = tr[lid].lcol;
		tr[id].rcol = tr[rid].rcol;
	}
	inline void pushdown(int id){
		if(tr[id].l == tr[id].r)	return;
		if(tr[id].cov){
			int t = tr[id].cov;
			tr[lid].cov = tr[rid].cov = t;
			tr[lid].lcol = tr[lid].rcol = t;
			tr[rid].lcol = tr[rid].rcol = t;
			tr[lid].cnt = tr[rid].cnt = 1;
			tr[id].cov = 0;
		}
	}
	void build(int id, int l, int r){
		tr[id].l = l, tr[id].r = r;
		if(tr[id].l == tr[id].r){
			tr[id].cnt = 1;
			tr[id].lcol = tr[id].rcol = a[fun[l]];
			return;
		}
		build(lid, l, mid);
		build(rid, mid+1, r);
		pushup(id);
	}
	void cover(int id, int l, int r, int c){
		pushdown(id);
		if(tr[id].l == l && tr[id].r == r){
			tr[id].cov = c;
			tr[id].cnt = 1;
			tr[id].lcol = tr[id].rcol = c;
			return;
		}
		if(r <= mid)	cover(lid, l, r, c);
		else if(l > mid)	cover(rid, l, r, c);
		else	cover(lid, l, mid, c), cover(rid, mid+1, r, c);
		pushup(id);
	}
	segTree query(int id, int l, int r){
		pushdown(id);
		if(tr[id].l == l && tr[id].r == r)	return tr[id];
		if(r <= mid)	return query(lid, l, r);
		else if(l > mid)	return query(rid, l, r);
		else{
			segTree t1 = query(lid, l, mid), t2 = query(rid, mid+1, r), t;
			t.cnt = t1.cnt + t2.cnt - (t1.rcol == t2.lcol);
			t.lcol = t1.lcol, t.rcol = t2.rcol;
			return t;
		}
	}
}seg;

void cover(int u, int v, int c){
	while(belong[u] != belong[v]){
		if(dep[belong[u]] < dep[belong[v]])	swap(u, v);
		seg.cover(1, st[belong[u]], st[u], c);
		u = fa[belong[u]];
	}
	if(dep[u] > dep[v])	swap(u, v);
	seg.cover(1, st[u], st[v], c);
}
int query(int u, int v){
	int ans = 0, lastu = 0, lastv = 0;
	while(belong[u] != belong[v]){
		if(dep[belong[u]] < dep[belong[v]])	swap(u, v), swap(lastu, lastv);
		segTree t = seg.query(1, st[belong[u]], st[u]);
		ans += t.cnt - (t.rcol == lastu);
		lastu = t.lcol;
		u = fa[belong[u]];
	}
	if(dep[u] > dep[v])	swap(u, v), swap(lastu, lastv);
	segTree t = seg.query(1, st[u], st[v]);
	ans += t.cnt - (t.lcol == lastu) - (t.rcol == lastv);
	return ans;
}

int main(){
	scanf("%d%d", &n, &m);
	for(int i = 1; i <= n; i++)
		scanf("%d", &a[i]), a[i]++;
	for(int i = 1; i < n; i++){
		scanf("%d%d", &u, &v);
		addEdge(u, v);
		addEdge(v, u);
	}
	dfs1(1, 1, 1);
	dfs2(1, 1);
	seg.build(1, 1, n);
	while(m--){
		scanf("%s", opt);
		if(opt[0] == 'C'){
			scanf("%d%d%d", &qa, &qb, &qc);
			qc++;
			cover(qa, qb, qc);
		}
		else if(opt[0] == 'Q'){
			scanf("%d%d", &qa, &qb);
			printf("%d\n", query(qa, qb));
		}
	}
	return 0;
}