P2486 [SDOI2011]染色 ·树链剖分+线段树

题解

我们知道树链剖分板子里第二次dfs将会为节点重新分配编号，每一条重链上的编号都是连续的，对应线段树的一段区间，编号越小越靠近树根，

在这道题里，主要难在大区间分割小区间时，小区间与小区间之间的端点的颜色，如果相同，这两个区间本该是连在一起的，显然，我们应该记下每个区间左右端点的颜色，

值得注意的是，在树链上统计答案的时候，由于树链的编号不一定连续，所以每统计一个区间的个数后，还需要和上一次的区间的端点进行比较，记作last_l,last_r，而在过程中，如果因为深度发生swap(u,v)，相当于整个链翻转，所以last_l和last_r也是需要交换的，

最后，如果查询的树链是人字形的，显然是需要判断两个端点的和上一次两个区间的端点状态

如果是直链，显然其中一端的上一次区间端点的状态永远处在初始值，就算判断两端也不会有影响

---

---

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 1e6 + 10;
const int INF = 0x3f3f3f3f;
int n, m, K;

string op;
namespace chain {//树链剖分板子
    int a[N];//实际点权

    struct egde {
        int to, next;
    } e[N];
    int head[N], tot;

    void add(int u, int v) {//一次建双边
        e[++tot] = {v, head[u]};
        head[u] = tot;
        e[++tot] = {u, head[v]};
        head[v] = tot;
    }

    int dep[N];//深度
    int f[N];//父节点
    int son[N];//重儿子
    int sz[N];//子树大小

    void dfs1(int u, int fa) {
        f[u] = fa;
        dep[u] = dep[fa] + 1;
        sz[u] = 1;
        for (int i = head[u]; i; i = e[i].next) {
            int v = e[i].to;
            if (v != fa) {
                dfs1(v, u);
                sz[u] += sz[v];
                if (sz[v] > sz[son[u]])
                    son[u] = v;
            }
        }
    }

    int dfn = 0, id[N], top[N];
    int w[N];//树链剖分后的点权
    void dfs2(int u, int tp) {
        id[u] = ++dfn;
        top[u] = tp;
        w[dfn] = a[u];//可能需要修改
        if (son[u])dfs2(son[u], tp);
        for (int i = head[u]; i; i = e[i].next) {
            int v = e[i].to;
            if (v != f[u] && v != son[u]) {
                dfs2(v, v);
            }
        }
    }
}
using namespace chain;

int Rc, Lc;//链端左右的颜色
namespace segment_tree {//线段树板子
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1

    int sum[N << 2];//颜色块的个数
    int lc[N << 2];//区间左边的点的颜色
    int rc[N << 2];//区间右边的点的颜色
    int lazy[N << 2];

    void pushup(int rt) {
        sum[rt] = sum[rt << 1] + sum[rt << 1 | 1] - (rc[rt << 1] == lc[rt << 1 | 1]);
        lc[rt] = lc[rt << 1];
        rc[rt] = rc[rt << 1 | 1];
    }

    void pushdown(int rt) {
        if (lazy[rt]) {
            sum[rt << 1] = sum[rt << 1 | 1] = 1;
            lc[rt << 1] = lc[rt << 1 | 1] = rc[rt << 1] = rc[rt << 1 | 1] = lc[rt];
            lazy[rt << 1] = lazy[rt << 1 | 1] = 1;
            lazy[rt] = 0;
        }
    }

    void build(int l, int r, int rt) {
        //记得初始化
        lazy[rt] = sum[rt] = lc[rt] = rc[rt] = 0;
        if (l == r) {
            lc[rt] = rc[rt] = w[l];
            sum[rt] = 1;
            return;
        }
        int mid = l + r >> 1;
        build(lson);
        build(rson);
        pushup(rt);
    }

    void update(int L, int R, int c, int l, int r, int rt) {
        if (L <= l && r <= R) {
            lc[rt] = rc[rt] = c;
            sum[rt] = 1;
            lazy[rt] = 1;
            return;
        }
        int mid = l + r >> 1;
        pushdown(rt);
        if (L <= mid) update(L, R, c, lson);
        if (R > mid) update(L, R, c, rson);
        pushup(rt);
    }

    int query(int L, int R, int x, int y, int l, int r, int rt) {
        if (x == l)Lc = lc[rt];//找大区间的两个端点的颜色
        if (y == r)Rc = rc[rt];

        if (L <= l && r <= R) {
            return sum[rt];
        }
        int mid = l + r >> 1;
        pushdown(rt);
        int res = 0;
        if (R <= mid) res = query(L, R, x, y, lson);
        else if (L > mid)res = query(L, R, x, y, rson);
        else res = query(L, R, x, y, lson) + query(L, R, x, y, rson) - (rc[rt << 1] == lc[rt << 1 | 1]);
        //内伤 这里写错了没看出来 卡了我一个晚上到自闭   rc[mid] == lc[mid+1] 这样写是错的!!!QAQ
        pushup(rt);
        return res;
    }
}
using namespace segment_tree;


void updRange(int u, int v, int c) {
    while (top[u] != top[v]) {
        if (dep[top[u]] < dep[top[v]])swap(u, v);
        update(id[top[u]], id[u], c, 1, n, 1);
        u = f[top[u]];
    }
    if (dep[u] > dep[v])swap(u, v);
    update(id[u], id[v], c, 1, n, 1);
}


int qRange(int u, int v) {
    int res = 0;
    int last_l = -1, last_r = -1;//上一次左端右端的颜色
    while (top[u] != top[v]) {
        if (dep[top[u]] < dep[top[v]])swap(u, v), swap(last_l, last_r);
        res += query(id[top[u]], id[u], id[top[u]], id[u], 1, n, 1);
        if (Rc == last_l) res--;
        u = f[top[u]];
        last_l = Lc;
    }
    if (dep[u] > dep[v])swap(u, v), swap(last_l, last_r);
    res += query(id[u], id[v], id[u], id[v], 1, n, 1);
    //重链要判断双边
    if (Lc == last_l) res--;
    if (Rc == last_r) res--;
    return res;
}

int main() {
    ios::sync_with_stdio(0);
    cin >> n >> m;
    for (int i = 1; i <= n; ++i) {
        cin >> a[i];
    }
    for (int i = 1, u, v; i < n; ++i) {
        cin >> u >> v;
        add(u, v);
    }
    dfs1(1, 0);
    dfs2(1, 1);
    build(1, n, 1);

    for (int i = 1, u, v, w; i <= m; ++i) {
        cin >> op >> u >> v;
        if (op == "C") {
            cin >> w;
            updRange(u, v, w);
        } else {
            cout << qRange(u, v) << endl;
        }
    }
    return 0;
}