题目
题目描述
给定一棵有n个节点的无根树和m个操作，操作有2类：
1、将节点a到节点b路径上所有点都染成颜色c；
2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），如“112221” 由3段组成：“11” 、“222” 和“1” 。
请你写一个程序依次完成这m个操作。
输入描述
第一行包含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）路径上的颜色段数量。
输出描述
对于每个询问操作，输出一行答案。
示例1
输入
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
输出
3
1
2
备注
对于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，保证给出的图是一棵树。除原数据外，还存在一组不计分的 hack 数据。
题解
知识点：树链剖分，线段树。
与路径查询修改有关，显然需要树链剖分。
颜色段问题可以用懒标记线段树简单处理。注意如果合并时，两段线段的合并处颜色相同，那么段数为段数和减 \(1\) ，否则为段数和。
其中，一个关键是使用树剖查询时，需要注意合并方向，因为颜色段的左右端点是需要区别的。因此，我们可以开一个数组分别记录左右的答案。为了配合查询时的 \(u,v\) 交换操作，采用一个变量 \(pos\) 表示当前 \(u\) 处在初始的左还是右， \(ans[pos]\) 就表示 \(u\) 那一段的答案。最后一次合并时，涉及三段的合并，注意其中一段是需要反转的，将端点颜色交换即可。
时间复杂度 \(O(n \log n + m \log ^2 n)\)
空间复杂度 \(O(n)\)
代码
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
struct HLD {
vector<int> siz, dep, fat, son, top, dfn, L, R;
HLD() {}
HLD(int rt, const vector<vector<int>> &g) { init(rt, g); }
void init(int rt, const vector<vector<int>> &g) {
assert(g.size() > rt);
int n = g.size() - 1;
siz.assign(n + 1, 0);
dep.assign(n + 1, 0);
fat.assign(n + 1, 0);
son.assign(n + 1, 0);
top.assign(n + 1, 0);
dfn.assign(n + 1, 0);
L.assign(n + 1, 0);
R.assign(n + 1, 0);
function<void(int, int)> dfsA = [&](int u, int fa) {
siz[u] = 1;
dep[u] = dep[fa] + 1;
fat[u] = fa;
for (auto v : g[u]) {
if (v == fa) continue;
dfsA(v, u);
siz[u] += siz[v];
if (siz[v] > siz[son[u]]) son[u] = v;
}
};
dfsA(rt, 0);
int dfncnt = 0;
function<void(int, int)> dfsB = [&](int u, int tp) {
top[u] = tp;
dfn[++dfncnt] = u;
L[u] = dfncnt;
if (son[u]) dfsB(son[u], tp);
for (auto v : g[u]) {
if (v == fat[u] || v == son[u]) continue;
dfsB(v, v);
}
R[u] = dfncnt;
};
dfsB(rt, rt);
}
};
template<class T, class F>
class SegmentTreeLazy {
int n;
vector<T> node;
vector<F> lazy;
void push_down(int rt) {
node[rt << 1] = lazy[rt](node[rt << 1]);
lazy[rt << 1] = lazy[rt](lazy[rt << 1]);
node[rt << 1 | 1] = lazy[rt](node[rt << 1 | 1]);
lazy[rt << 1 | 1] = lazy[rt](lazy[rt << 1 | 1]);
lazy[rt] = F();
}
void update(int rt, int l, int r, int x, int y, F f) {
if (r < x || y < l) return;
if (x <= l && r <= y) return node[rt] = f(node[rt]), lazy[rt] = f(lazy[rt]), void();
push_down(rt);
int mid = l + r >> 1;
update(rt << 1, l, mid, x, y, f);
update(rt << 1 | 1, mid + 1, r, x, y, f);
node[rt] = node[rt << 1] + node[rt << 1 | 1];
}
T query(int rt, int l, int r, int x, int y) {
if (r < x || y < l) return T();
if (x <= l && r <= y) return node[rt];
push_down(rt);
int mid = l + r >> 1;
return query(rt << 1, l, mid, x, y) + query(rt << 1 | 1, mid + 1, r, x, y);
}
public:
SegmentTreeLazy(int _n = 0) { init(_n); }
SegmentTreeLazy(const vector<T> &src) { init(src); }
void init(int _n) {
n = _n;
node.assign(n << 2, T());
lazy.assign(n << 2, F());
}
void init(const vector<T> &src) {
assert(src.size() >= 2);
init(src.size() - 1);
function<void(int, int, int)> build = [&](int rt, int l, int r) {
if (l == r) return node[rt] = src[l], void();
int mid = l + r >> 1;
build(rt << 1, l, mid);
build(rt << 1 | 1, mid + 1, r);
node[rt] = node[rt << 1] + node[rt << 1 | 1];
};
build(1, 1, n);
}
void update(int x, int y, F f) { update(1, 1, n, x, y, f); }
T query(int x, int y) { return query(1, 1, n, x, y); }
};
struct T {
int lc = 0, rc = 0;
int cnt = 0;
friend T operator+(const T &a, const T &b) {
if (!a.cnt) return b;
if (!b.cnt) return a;
return {
a.lc, b.rc,
a.cnt + b.cnt - (a.rc == b.lc)
};
}
};
struct F {
int upd;
T operator()(const T &x) {
if (!upd) return x;
return {
upd, upd,
1,
};
}
F operator()(const F &g) {
if (!upd) return g;
return { upd };
}
};
const int N = 100007;
int c[N];
vector<int> g[N];
HLD hld;
SegmentTreeLazy<T, F> sgt;
void path_update(int u, int v, int w) {
auto &top = hld.top;
auto &dep = hld.dep;
auto &L = hld.L;
auto &fat = hld.fat;
while (top[u] != top[v]) {
if (dep[top[u]] < dep[top[v]]) swap(u, v);
sgt.update(L[top[u]], L[u], { w });
u = fat[top[u]];
}
if (dep[u] > dep[v]) swap(u, v);
sgt.update(L[u], L[v], { w });
}
int path_query(int u, int v) {
auto &top = hld.top;
auto &dep = hld.dep;
auto &L = hld.L;
auto &fat = hld.fat;
bool pos = 0;
T ans[2] = { T(),T() };
while (top[u] != top[v]) {
if (dep[top[u]] < dep[top[v]]) swap(u, v), pos ^= 1;
ans[pos] = sgt.query(L[top[u]], L[u]) + ans[pos];
u = fat[top[u]];
}
if (dep[u] > dep[v]) swap(u, v), pos ^= 1;
swap(ans[pos].lc, ans[pos].rc);
return (ans[pos] + sgt.query(L[u], L[v]) + ans[pos ^ 1]).cnt;
}
int main() {
std::ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
int n, m;
cin >> n >> m;
for (int i = 1;i <= n;i++) cin >> c[i];
for (int i = 1;i <= n - 1;i++) {
int u, v;
cin >> u >> v;
g[u].push_back(v);
g[v].push_back(u);
}
hld.init(1, vector<vector<int>>(g, g + n + 1));
vector<T> c_src(n + 1);
for (int i = 1;i <= n;i++) c_src[hld.L[i]] = { c[i],c[i],1 };
sgt.init(c_src);
while (m--) {
char op;
int a, b;
cin >> op >> a >> b;
if (op == 'C') {
int c;
cin >> c;
path_update(a, b, c);
}
else cout << path_query(a, b) << '\n';
}
return 0;
}