HYSBZ - 2243 染色(树剖 ＋ 区间合并)

题目大意：给定一棵有n个节点的无根树和m个操作，操作有2类：
1、将节点a到节点b路径上所有点都染成颜色c；
2、询问节点a到节点b路径上的颜色段数量（连续相同颜色被认为是同一段），如“112221”由3段组成：“11”、“222”和“1”。
请你写一个程序依次完成这m个操作。

解题思路：维护三个值，sum表示该区间的颜色段数量，left表示该区间的最左边的颜色，right表示该区间最右边的颜色
区间合并的时候，判断一下两个区间的合并处颜色是否相同，如果相同，数量减一
最后是询问操作，再模拟一次回溯路径，然后判断一下两个区间的相交之处颜色是否相同，如果相同，答案减一

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;
const int N = 100010;
const int M = 100010 << 2;

struct Edge{
    int u, v, next;
    Edge() {}
    Edge(int u, int v, int next): u(u), v(v), next(next) {}
}E[N * 2];
int top[N], size[N], id[N], dep[N], son[N], f[N], head[N], val[N], color[N];
int left[M], right[M], sum[M], setv[M];
int n, m, tot, idx;

void dfs1(int u, int fa, int depth) {
    dep[u] = depth; size[u] = 1; son[u] = 0; f[u] = fa;
    for (int i = head[u]; ~i; i = E[i].next) {
        int v = E[i].v;
        if (v == fa) continue;
        dfs1(v, u, depth + 1);
        size[u] += size[v];
        if (size[v] > size[son[u]]) son[u] = v;
    }
}

void dfs2(int u, int tp) {
    top[u] = tp;
    id[u] = ++idx;
    if (son[u]) dfs2(son[u], tp);

    for (int i = head[u]; ~i; i = E[i].next) {
        int v = E[i].v;
        if (v == f[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}

void PushUp(int u, int l, int r) {
    left[u] = left[u << 1];
    right[u] = right[u << 1 | 1];
    sum[u] = sum[u << 1] + sum[u << 1 | 1];
    if (right[u << 1] == left[u << 1 | 1]) sum[u]--;
}

void build(int u, int l, int r) {
    setv[u] = -1;
    if (l == r) {
        sum[u] = 1;
        left[u] = right[u] = color[l];
        return ;
    }
    int mid = (l + r) >> 1;
    build(u << 1, l, mid);
    build(u << 1 | 1, mid + 1, r);
    PushUp(u, l, r);
}

void PushDown(int u, int l, int r) {
    setv[u << 1] = setv[u << 1 | 1] = setv[u];
    left[u << 1] = right[u << 1] = setv[u];
    left[u << 1 | 1] = right[u << 1 | 1] = setv[u];
    sum[u << 1] = sum[u << 1 | 1] = 1;
    setv[u] = -1;
}

void Modify(int u, int l, int r, int L, int R, int c) {
    if (l == L && r == R) {
        setv[u] = left[u] = right[u] = c;
        sum[u] = 1;
        return ;
    }
    if (~setv[u]) PushDown(u, l, r);
    int mid = (l + r) >> 1;
    if (R <= mid) Modify(u << 1, l, mid, L, R, c);
    else if (L > mid) Modify(u << 1 | 1, mid + 1, r, L, R, c);
    else {
        Modify(u << 1, l, mid, L, mid, c);
        Modify(u << 1 | 1, mid + 1, r, mid + 1, R, c);
    }
    PushUp(u, l, r);
}

int Query(int u, int l, int r, int L, int R) {
    if (l == L && r == R) return sum[u];
    if (~setv[u]) PushDown(u, l, r);

    int mid = (l + r) >> 1;
    if (R <= mid) return Query(u << 1, l, mid, L, R);
    else if (L > mid) return Query(u << 1 | 1, mid + 1, r, L, R);
    else {
        int t1 = Query(u << 1, l, mid, L, mid);
        int t2 = Query(u << 1 | 1, mid + 1, r, mid + 1, R);
        if (right[u << 1] == left[u << 1 | 1]) return t1 + t2 - 1;
        else return t1 + t2;
    }
}

int FindColor(int u, int l, int r, int pos) {
    if (l == r) return left[u];
    if (~setv[u]) PushDown(u, l, r);
    int mid = (l + r) >> 1;
    if (pos <= mid) return FindColor(u << 1, l, mid, pos);
    else return FindColor(u << 1 | 1, mid + 1, r, pos);
}

void gao(int u, int v, int c) {
    int t1 = u, t2 = v;
    int tp1 = top[u], tp2 = top[v];
    int ans = 0;
    while (tp1 != tp2) {
        if (dep[tp1] < dep[tp2]) {
            swap(tp1, tp2);
            swap(u, v);
        }
        if (c == -1) ans += Query(1, 1, idx, id[tp1], id[u]);
        else Modify(1, 1, idx, id[tp1], id[u], c);
        u = f[tp1];
        tp1 = top[u];
    }
    if (dep[u] > dep[v]) swap(u, v);
    if (c == -1) ans += Query(1, 1, idx, id[u], id[v]);
    else Modify(1, 1, idx, id[u], id[v], c);
    if (c == -1) {
        if (top[t1] == top[t2]) {
            printf("%d\n", ans);
            return ;
        }
        int tp3 = top[t1];
        while (tp3 != tp1) {
            int fa = f[tp3];
            int c1 = FindColor(1, 1, idx, id[tp3]);
            int c2 = FindColor(1, 1, idx, id[fa]);
            if (c1 == c2) ans--;
            tp3 = top[fa];
        }

        int tp4 = top[t2];
        while (tp4 != tp1) {
            int fa = f[tp4];
            int c1 = FindColor(1, 1, idx, id[tp4]);
            int c2 = FindColor(1, 1, idx, id[fa]);
            if (c1 == c2) ans--;
            tp4 = top[fa];
        }
        printf("%d\n", ans);
    }
}

void solve() {
    char op[4];
    int a, b, c;
    while (m--) {
        scanf("%s", op);
        if (op[0] == 'Q') {
            scanf("%d%d", &a, &b);
            gao(a, b, -1);
        }
        else {
            scanf("%d%d%d", &a, &b, &c);
            gao(a, b, c);
        }
    }
}

void AddEdge(int u, int v) {
    E[tot] = Edge(u, v, head[u]);
    head[u] = tot++;
    E[tot] = Edge(v, u, head[v]);
    head[v] = tot++;
}

void init() {
    memset(head, -1, sizeof(head));
    tot = 0;

    int u, v;
    for (int i = 1; i <= n; i++)
        scanf("%d", &val[i]);

    for (int i = 1; i < n; i++) {
        scanf("%d%d", &u, &v);
        AddEdge(u, v);
    }

    idx = 0;
    dfs1(1, -1, 1);
    dfs2(1, 1);

    for (int i = 1; i <= n; i++)
        color[id[i]] = val[i];
    build(1, 1, idx);
}

int main() {
    while (scanf("%d%d", &n, &m) != EOF) {  
        init();
        solve();
    }
    return 0;
}