洛谷P2486 [SDOI2011]染色 题解 树链剖分+线段树

题目链接：https://www.luogu.org/problem/P2486
首先这是一道树链剖分+线段树的题。
线段树部分和 codedecision P1112 区间连续段 一模一样，所以我们在做这道题目之前最好去做一下这道题目的练习。
然后就是树链剖分的部分。
此部分支持两种操作：

• 更新：这部分比较好实现；
• 查询：这部分需要你记录树链查询的时候的每一条边的信息，然后将这些信息进行汇总，处理起来稍有一些繁琐。

实现代码如下：

#include <bits/stdc++.h>
using namespace std;
#define INF (1<<29)
const int maxn = 100010;
int fa[maxn],
    dep[maxn],
    size[maxn],
    son[maxn],
    top[maxn],
    seg[maxn], seg_cnt,
    rev[maxn];
vector<int> g[maxn];
void dfs1(int u, int p) {
    size[u] = 1;
    for (vector<int>::iterator it = g[u].begin(); it != g[u].end(); it ++) {
        int v = (*it);
        if (v == p) continue;
        fa[v] = u;
        dep[v] = dep[u] + 1;
        dfs1(v, u);
        size[u] += size[v];
        if (size[v] >size[son[u]]) son[u] = v;
    }
}
void dfs2(int u, int tp) {
    seg[u] = ++seg_cnt;
    rev[seg_cnt] = u;
    top[u] = tp;
    if (son[u]) dfs2(son[u], tp);
    for (vector<int>::iterator it = g[u].begin(); it != g[u].end(); it ++) {
        int v = (*it);
        if (v == fa[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}
struct Node {
    int l, r, cnt;
    Node () {}
    Node (int _l, int _r, int _cnt) { l = _l; r = _r; cnt = _cnt; }
    Node reverse() { return Node(r, l, cnt); }
} tree[maxn<<2];
int n, lazy[maxn<<2], init_color[maxn];
#define lson l, mid, rt<<1
#define rson mid+1, r, rt<<1|1
void push_up(int rt) {
    tree[rt].l = tree[rt<<1].l;
    tree[rt].r = tree[rt<<1|1].r;
    tree[rt].cnt = tree[rt<<1].cnt + tree[rt<<1|1].cnt - (tree[rt<<1].r == tree[rt<<1|1].l ? 1 : 0);
}
void push_down(int rt) {
    if (lazy[rt]) {
        lazy[rt<<1] = lazy[rt<<1|1] = lazy[rt];
        tree[rt<<1].cnt = tree[rt<<1|1].cnt = 1;
        tree[rt<<1].l = tree[rt<<1].r = tree[rt<<1|1].l = tree[rt<<1|1].r = lazy[rt];
        lazy[rt] = 0;
    }
}
void build(int l, int r, int rt) {
    if (l == r) {
        tree[rt] = Node(init_color[rev[l]], init_color[rev[l]], 1);
        return;
    }
    int mid = (l + r) / 2;
    build(lson);
    build(rson);
    push_up(rt);
}
void update(int L, int R, int v, int l, int r, int rt) {
    if (L <= l && r <= R) {
        tree[rt] = Node(v, v, 1);
        lazy[rt] = v;
        return;
    }
    push_down(rt);
    int mid = (l + r) / 2;
    if (L <= mid) update(L, R, v, lson);
    if (R > mid) update(L, R, v, rson);
    push_up(rt);
}
Node query(int L, int R, int l, int r, int rt) {
    if (L <= l && r <= R) return tree[rt];
    push_down(rt);
    int mid = (l + r) / 2;
    if (L > mid) return query(L, R, rson);
    else if (R <= mid) return query(L, R, lson);
    else {
        Node a = query(L, R, lson);
        Node b = query(L, R, rson);
        return Node(a.l, b.r, a.cnt + b.cnt - (a.r == b.l ? 1 : 0));
    }
}

void chain_update(int u, int v, int val) {
    while (top[u] != top[v]) {
        if (dep[top[u]] < dep[top[v]]) swap(u, v);
        update(seg[top[u]], seg[u], val, 1, n, 1);
        u = fa[top[u]];
    }
    if (dep[u] < dep[v]) swap(u, v);
    update(seg[v], seg[u], val, 1, n, 1);
}
vector<Node> res1, res2, res;
void chain_query(int u, int v) {
    res1.clear();
    res2.clear();
    res.clear();
    while (top[u] != top[v]) {
        if (dep[top[u]] > dep[top[v]]) {
            res1.push_back(query(seg[top[u]], seg[u], 1, n, 1));
            u = fa[top[u]];
        }
        else {
            res2.push_back(query(seg[top[v]], seg[v], 1, n, 1));
            v = fa[top[v]];
        }
    }
    if (dep[u] > dep[v]) res1.push_back(query(seg[v], seg[u], 1, n, 1));
    else res2.push_back(query(seg[u], seg[v], 1, n, 1));
    int sz = res1.size();
    for (int i = 0; i < sz; i ++) res.push_back(res1[i].reverse());
    sz = res2.size();
    for (int i = sz-1; i >= 0; i --) res.push_back(res2[i]);
    Node tmp = res[0];
    sz = res.size();
    for (int i = 1; i < sz; i ++) {
        int delta = (tmp.r == res[i].l);
        tmp.cnt += res[i].cnt - delta;
        tmp.r = res[i].r;
    }
    cout << tmp.cnt << endl;
}

int m, a, b, c;
char op[2];
int main() {
    cin >> n >> m;
    for (int i = 1; i <= n; i ++) cin >> init_color[i];
    for (int i = 1; i < n; i ++) {
        int u, v;
        cin >> u >>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dep[1] = fa[1] = 1;
    dfs1(1, -1);
    dfs2(1, 1);
    build(1, n, 1);
    while (m --) {
        cin >> op;
        if (op[0] == 'C') {
            cin >> a >> b >> c;
            chain_update(a, b, c);
        }
        else {
            cin >> a >> b;
            chain_query(a, b);
        }
    }
    return 0;
}