P2486 染色 (树链剖分)

题目链接: P2486 染色

2021.11.5更新了树链信息合并的写法

大致题意

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

1. 将节点$a$到节点$b$径上的所有点(包括$a$和$b$) 都染成颜色$c$。
2. 询问节点$a$到节点$b$的路径上的颜色段数量。颜色段的定义是极长的连续相同颜色被认为是一段。例如 112221 由三段组成：11、222、1。

解题思路

树链剖分

首先我们考虑如何在一段序列上实现上述操作.
我们发现只需要线段树 + 区间合并即可实现.

但本题是树上操作, 我们考虑采用树剖. 把树上路径转化为区间进行操作.

考虑到修改操作, 我们直接套板子进行区间修改即可.

考虑到查询操作, 由于树剖把一条路径拆成了几条树链, 因此我们需要考虑两个区间相互拼接的问题.
我们每次$ask(id[top[a]],id[a])$的颜色, 然后使$a=p[top[a]]$, 那么相连点为$top[a]$和$p[top[a]]$. 如果这两个点的颜色相同, 我们便对答案$-1$.

---

当然, 对于查询操作我们也可以考虑直接去合并树链信息.

➡️树链信息合并参考博客⬅️

在下文中会给出两种做法的代码.

AC代码

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 1; i <= (n); ++i)
using namespace std;
typedef long long ll;
const int N = 1E5 + 10;
int w[N];
vector<int> edge[N];

int p[N], dep[N], sz[N], son[N];
void dfs1(int x = 1, int fa = 0) {
	p[x] = fa, dep[x] = dep[fa] + 1, sz[x] = 1;
	for (auto& to : edge[x]) {
		if (to == fa) continue;
		dfs1(to, x);
		sz[x] += sz[to];
		if (sz[to] > sz[son[x]]) son[x] = to;
	}
}
int id[N], nw[N], top[N], ind;
void dfs2(int x = 1, int tp = 1) {
	id[x] = ++ind, nw[ind] = w[x], top[x] = tp;

	if (!son[x]) return;
	dfs2(son[x], tp);

	for (auto& to : edge[x]) {
		if (to == p[x] or to == son[x]) continue;
		dfs2(to, to);
	}
}


struct node {
	int l, r;
	int lc, rc, num;
	int col;
}t[N << 2];
void pushdown(node& op, int col) {
	op.lc = op.rc = col; op.num = 1;
	op.col = col;
}
void pushdown(int x) {
	if (!t[x].col) return;
	pushdown(t[x << 1], t[x].col), pushdown(t[x << 1 | 1], t[x].col);
	t[x].col = 0;
}

void pushup(node& p, const node& l, const node& r) {
	p.lc = l.lc, p.rc = r.rc;
	p.num = l.num + r.num - (l.rc == r.lc);
}
void pushup(int x) { pushup(t[x], t[x << 1], t[x << 1 | 1]); }

void build(int l, int r, int x = 1) {
	t[x] = { l, r, nw[l], nw[l], 1, 0 };
	if (l == r) return;
	int mid = l + r >> 1;
	build(l, mid, x << 1), build(mid + 1, r, x << 1 | 1);
	pushup(x);
}

void modify(int l, int r, int c, int x = 1) {
	if (l <= t[x].l and r >= t[x].r) {
		pushdown(t[x], c);
		return;
	}
	pushdown(x);
	int mid = t[x].l + t[x].r >> 1;
	if (l <= mid) modify(l, r, c, x << 1);
	if (r > mid) modify(l, r, c, x << 1 | 1);
	pushup(x);
}

node ask(int l, int r, int x = 1) {
	if (l <= t[x].l and r >= t[x].r) return t[x];
	pushdown(x);
	int mid = t[x].l + t[x].r >> 1;
	if (r <= mid) return ask(l, r, x << 1);
	if (l > mid) return ask(l, r, x << 1 | 1);
	node left = ask(l, r, x << 1), right = ask(l, r, x << 1 | 1);
	node res = { left.l, right.r, 0, 0, 0, 0 };
	pushup(res, left, right);
	return res;
}


void modify_route(int a, int b, int c) {
	while (top[a] != top[b]) {
		if (dep[top[a]] < dep[top[b]]) swap(a, b);
		modify(id[top[a]], id[a], c);
		a = p[top[a]];
	}
	if (dep[a] > dep[b]) swap(a, b);
	modify(id[a], id[b], c);
}

namespace 查询操作的两种写法{ // 代码使用时, 把命名空间的代码复制出去即可.
    /* 写法一: 判断两端点颜色是否相同 */
    int ask_route(int a, int b) {
        int res = 0;
        while (top[a] != top[b]) {
            if (dep[top[a]] < dep[top[b]]) swap(a, b);
            auto qaq = ask(id[top[a]], id[a]);
            a = p[top[a]];
            res += qaq.num - (qaq.lc == ask(id[a], id[a]).lc);
        }
        if (dep[a] > dep[b]) swap(a, b);
        res += ask(id[a], id[b]).num;

        return res;
    }
    
    
    /* 写法二: 树链信息合并 */
    node merge(const node& l, const node& r) {
        node res = { 0, 0, 0, 0, 0, 0 };
        pushup(res, l, r);
        return res;
	}
    int ask_route(int a, int b) {
        node la = { 0, 0, 0, 0, 0, 0 }, lb = { 0, 0, 0, 0, 0, 0 };
        while (top[a] != top[b]) {
            if (dep[top[a]] < dep[top[b]]) swap(a, b), swap(la, lb);
            node qaq = ask(id[top[a]], id[a]);
            la = merge(qaq, la);

            a = p[top[a]];
        }

        if (id[a] > id[b]) swap(a, b), swap(la, lb);
        node qaq = ask(id[a], id[b]);
        lb = merge(qaq, lb);
        swap(la.lc, la.rc);
        return merge(la, lb).num;
    }
}
int main()
{
	int n, m; cin >> n >> m;
	rep(i, n) scanf("%d", &w[i]);
	rep(i, n - 1) {
		int a, b; scanf("%d %d", &a, &b);
		edge[a].push_back(b), edge[b].push_back(a);
	}

	dfs1(), dfs2();
	build(1, n);

	rep(i, m) {
		char s[2]; scanf("%s", s);
		if (*s == 'C') {
			int a, b, c; scanf("%d %d %d", &a, &b, &c);
			modify_route(a, b, c);
		}
		else {
			int a, b; scanf("%d %d", &a, &b);
			printf("%d\n", ask_route(a, b));
		}
	}

	return 0;
}

END