題目鏈接

題解

比較裸的樹鏈剖分 好像樹鏈剖分的題都很裸

線段樹中維護一個區間最左和最右的顏色，和答案
合並判斷一下中間一段就可以了

比較考驗代碼能力

Code

#include<bits/stdc++.h>
#define LL long long
#define RG register
using namespace std;

inline int gi() {
    int f = 1, s = 0;
    char c = getchar();
    while (c != '-' && (c < '0' || c > '9')) c = getchar();
    if (c == '-') f = -1, c = getchar();
    while (c >= '0' && c <= '9') s = s*10+c-'0', c = getchar();
    return f == 1 ? s : -s;
}
const int N = 100010;
struct node {
    int to, next;
}g[N<<1];
int last[N], gl;
inline void add(int x, int y) {
    g[++gl] = (node) {y, last[x]};
    last[x] = gl;
    g[++gl] = (node) {x, last[y]};
    last[y] = gl;
    return ;
}

int c[N];

struct Segment_tree {
    int l, r, v, lazy;
}t[N<<2];

int fa[N], siz[N], son[N], cnt, s[N], top[N], id[N], dep[N];
void dfs1(int u, int f) {
    siz[u] = 1; fa[u] = f;
    int MAX = 0;
    for (int i = last[u]; i; i = g[i].next) {
        int v = g[i].to;
        if (v == f) continue;
        dep[v] = dep[u]+1;
        dfs1(v, u);
        siz[u] += siz[v];
        if (MAX < siz[v]) MAX = siz[v], son[u] = v;
    }
    return ;
}

void dfs2(int u, int topf) {
    top[u] = topf;
    s[++cnt] = c[u];
    id[u] = cnt;
    if (!son[u]) return ;
    dfs2(son[u], topf);
    for (int i = last[u]; i; i = g[i].next) {
        int v = g[i].to;
        if (v == fa[u] || v == son[u]) continue;
            dfs2(v, v);
    }
    return ;
}
#define mid ((l+r)>>1)
#define ls (rt<<1)
#define rs (rt<<1|1)


void pushup(int rt) {   
    t[rt].l = t[ls].l; t[rt].r = t[rs].r;
    t[rt].v = t[ls].v + t[rs].v;
    if (t[ls].r == t[rs].l) t[rt].v--;
    return ;
}

void build(int rt, int l, int r) {
    if (l == r) {
        t[rt] = (Segment_tree) {s[l], s[l], 1, 0};
        return ;
    }
    build(ls, l, mid); build(rs, mid+1, r);
    pushup(rt);
    return ;
}
void pushdown(int rt) {
    int lazy = t[rt].lazy;
    t[rt].lazy = 0;
    if (lazy) {
        t[rs].lazy = t[ls].lazy = lazy;
        t[rs].l = t[rs].r = t[ls].l = t[ls].r = lazy;
        t[rs].v = t[ls].v = 1;
    }
    return ;
}
void update(int rt, int l, int r, int L, int R, int k) {
    if (L <= l && r <= R) {
        t[rt].v = 1;
        t[rt].lazy = t[rt].l = t[rt].r = k;
        return ;
    }
    pushdown(rt);
    if (L <= mid) update(ls, l, mid, L, R, k);
    if (R > mid) update(rs, mid+1, r, L, R, k);
    pushup(rt);
    return ;
}

struct zzy {
    int l, r, v;
}; 

zzy merge(zzy a, zzy b) {
    if (!a.v) return b;
    if (!b.v) return a;
    return (zzy) {a.l, b.r, a.v+b.v-(a.r == b.l)};
}

zzy query(int rt, int l, int r, int L, int R) {
    if (L <= l && r <= R)
        return (zzy) {t[rt].l, t[rt].r, t[rt].v};
    zzy z = (zzy) {0, 0, 0};
    pushdown(rt);
    if (L <= mid)
        z = query(ls, l, mid, L, R);
    if (R > mid)
        z = merge(z, query(rs, mid+1, r, L, R));
    return z;
}

void upway(int x, int y, int z) {
    while (top[x] != top[y]) {
        if (dep[top[x]] < dep[top[y]]) swap(x, y);
        update(1, 1, cnt, id[top[x]], id[x], z);
        x = fa[top[x]];
    }
    if (dep[x] > dep[y]) swap(x, y);
    update(1, 1, cnt, id[x], id[y], z);
    return ;
}

int qway(int x, int y) {
    zzy zx, zy;
    zx = zy = (zzy) {0, 0, 0};
    while (top[x] != top[y]) {
        if (dep[top[x]] >= dep[top[y]]) {
            zx = merge(query(1, 1, cnt, id[top[x]], id[x]), zx);
            x = fa[top[x]];
        }
        else {
            zy = merge(query(1, 1, cnt, id[top[y]], id[y]), zy);
            y = fa[top[y]];
        }
    }
    if (dep[x] > dep[y]) swap(x, y), swap(zx, zy);
    zy = merge(query(1, 1, cnt, id[x], id[y]), zy);
    zx.v += zy.v;
    if (zx.l == zy.l) zx.v--;
    return zx.v;
}

int main() {
    int n = gi(), m = gi();
    for (int i = 1; i <= n; i++) c[i] = gi();
    for (int i = 1; i < n; i++) add(gi(), gi());
    dfs1(1, 0);
    dfs2(1, 1);
    build(1, 1, n);
    while (m--) {
        char c; cin>>c;
        if (c == 'C') {
            int x = gi(), y = gi(), z = gi();
            upway(x, y, z);
        }
        else {
            int x = gi(), y = gi();
            printf("%d\n", qway(x, y));
        }
    }
    return 0;
}
 

洛谷 P2486 [SDOI2011]染色（樹鏈剖分+線段樹）