题目

题目链接

P2486 [SDOI2011] 染色

题目描述

给定一棵 $n$ 个节点的无根树，共有 $m$ 个操作，操作分为两种：

1. 将节点 $a$ 到节点 $b$ 的路径上的所有点（包括 $a$ 和 $b$）都染成颜色 $c$。
2. 询问节点 $a$ 到节点 $b$ 的路径上的颜色段数量。

颜色段的定义是极长的连续相同颜色被认为是一段。例如 112221 由三段组成：11、222、1。

输入格式

输入的第一行是用空格隔开的两个整数，分别代表树的节点个数 $n$ 和操作个数 $m$。

第二行有 $n$ 个用空格隔开的整数，第 $i$ 个整数 $w_i$ 代表结点 $i$ 的初始颜色。

第 $3$ 到第 $(n + 1)$ 行，每行两个用空格隔开的整数 $u, v$，代表树上存在一条连结节点 $u$ 和节点 $v$ 的边。

第 $(n + 2)$ 到第 $(n + m + 1)$ 行，每行描述一个操作，其格式为：

每行首先有一个字符 $op$，代表本次操作的类型。

• 若 $op$ 为 C，则代表本次操作是一次染色操作，在一个空格后有三个用空格隔开的整数 $a, b, c$，代表将 $a$ 到 $b$ 的路径上所有点都染成颜色 $c$。
• 若 $op$ 为 Q，则代表本次操作是一次查询操作，在一个空格后有两个用空格隔开的整数 $a, b$，表示查询 $a$ 到 $b$ 路径上的颜色段数量。

输出格式

对于每次查询操作，输出一行一个整数代表答案。

对于 $100\%$ 的数据，$1 \leq n, m \leq 10^5$，$1 \leq w_i, c \leq 10^9$，$1 \leq a, b, u, v \leq n$，$op$ 一定为 C 或 Q，保证给出的图是一棵树。

解析

思路

本题的思维难度很小，用线段树维护区间颜色段数，重点在如何实现去重。

现有的轻重链剖分题解，基本上都是开了空间大小为 4 倍的 $\mathtt{lc,rc}$ 数组，分别维护区间左右端点颜色，其实根本没有必要，

一个 1 倍空间的 $\mathtt{col}$ 数组足以解决问题，而且贼好写。

1. 线段树上合并时的去重（左儿子右端点 和 右儿子左端点）

我们用 $\mathtt{sum[i]}$ 表示 $i$ 所管辖区间内的颜色段数。若 $i$ 所管辖的区间为 $\mathtt{[l,r]}$，那么左儿子管辖区间为 $\mathtt{[l,mid]}$，右儿子管辖区间为 $\mathtt{[mid+1,r]}$，我们想要判断合并时会不会计重，其实只是想知道 $\mathtt{col[mid]}$ 是否等于 $\mathtt{col[mid+1]}$；

这样的好处是：上方的节点不用继承儿子的颜色，因为我们只要知道 $i$ 的管辖区间，就可以知道左右端点的颜色。

同样的，懒标记下传时，只需要将左儿子右端点 $\mathtt{col[mid]}$ 和右儿子左端点 $\mathtt{col[mid+1]}$ 修改就可以了。

2. 重链跳跃时的去重（重链连接处 即 轻链两端点）

大部分题解都通过记录上次跳跃位置实现，之所以要记录上次跳的位置，是因为由于有懒标记，只有查询后的颜色才是正确的（该下放的都下放了）。

其实有一种更好写的方法，我们可以在全部查完后，重新跳一遍来判断重链连接处，详见代码。

代码

#include <cstdio>
#include <vector>
using namespace std;
const int maxn = 100005;
int n, m, o[maxn], col[maxn];
vector<int> e[maxn];
int dep[maxn], fa[maxn], son[maxn], size[maxn];//dfs1
int dfn[maxn], top[maxn], clk;//clk:dfs_clock dfs2
struct SegTree {
	int sum[maxn<<2], tag[maxn<<2];
	#define ls i<<1
	#define rs i<<1|1
	inline void push_up(int i, int mid) { 
		sum[i] = sum[ls] + sum[rs];
		if(col[mid] == col[mid+1]) --sum[i];
	}
	inline void push_down(int i, int mid) {
		tag[ls] = tag[rs] = col[mid] = col[mid+1] = tag[i];
		sum[ls] = sum[rs] = 1;
		tag[i] = 0;
	}
	void build(int i, int l, int r) {
		if(l == r) { sum[i] = 1; return; }
		int mid = (l+r) >> 1;
		build(ls, l, mid);
		build(rs, mid+1, r);
		push_up(i, mid);
	}
	void ins(int i, int l, int r, int x, int y, int k) {
		if(x <= l and r <= y) { sum[i] = 1; col[l] = col[r] = tag[i] = k; return; }
		int mid = (l+r) >> 1;
		if(tag[i]) push_down(i, mid);
		if(x <= mid) ins(ls, l, mid, x, y, k);
		if(y > mid) ins(rs, mid+1, r, x, y, k);
		push_up(i, mid);
	}
	int query(int i, int l, int r, int x, int y) {
		if(x <= l and r <= y) return sum[i];
		int mid = (l+r) >> 1, ans = 0;
		if(tag[i]) push_down(i, mid);
		if(x <= mid) ans = query(ls, l, mid, x, y);
		if(y > mid) ans += query(rs, mid+1, r, x, y);
		if(x <= mid and y > mid and col[mid] == col[mid+1]) --ans;
		return ans;
	}
} t;
void dfs1(int u, int f) {
	fa[u] = f;
	dep[u] = dep[f] + 1;
	size[u] = 1;
	for(auto v : e[u])
		if(v != f) {
			dfs1(v, u);
			size[u] += size[v];
			if(size[v] > size[son[u]]) son[u] = v;
		} 
}
void dfs2(int u, int topf) {
	dfn[u] = ++clk;
	col[clk] = o[u];
	top[u] = topf;
	if(!son[u]) return;
	dfs2(son[u], topf);
	for(auto v : e[u])
		if(!dfn[v]) dfs2(v, v);
}
void modify(int x, int y, int k) {
	while(top[x] != top[y]) {
		if(dep[top[x]] < dep[top[y]]) swap(x, y);
		t.ins(1, 1, n, dfn[top[x]], dfn[x], k);
		x = fa[top[x]];
	}
	if(dfn[x] > dfn[y]) swap(x, y);
	t.ins(1, 1, n, dfn[x], dfn[y] ,k);
}
int inquire(int x, int y) {
	int u = x, v = y, ans = 0;//记录 x,y 因为还要跳第二次
	while(top[x] != top[y]) {
		if(dep[top[x]] < dep[top[y]]) swap(x, y);
		ans += t.query(1, 1, n, dfn[top[x]], dfn[x]);
		x = fa[top[x]];
	}
	if(dfn[x] > dfn[y]) swap(x, y);
	ans += t.query(1, 1, n, dfn[x], dfn[y]);
	while(top[u] != top[v]) {
		if(dep[top[u]] < dep[top[v]]) swap(u, v);
		if(col[dfn[top[u]]] == col[dfn[fa[top[u]]]]) --ans;
		u = fa[top[u]];//top[u] 和 fa[top[u]] 是轻链端点
	}
	return ans;
}
char c[2];
int main() {
	scanf("%d %d", &n, &m);
	for(int i = 1; i <= n; ++i) scanf("%d", &o[i]);
	for(int i = 1, u, v; i <= n; ++i)
		scanf("%d %d", &u, &v), e[u].push_back(v), e[v].push_back(u);
	dfs1(1, 1);
	dfs2(1, 1);
	t.build(1, 1, n);
	for(int i = 1, x, y, k; i <= m; ++i) {
		scanf("%s %d %d", c, &x, &y);
		if(c[0] == 'C') scanf("%d", &k), modify(x, y, k);
		else printf("%d\n", inquire(x, y));
	}
}