HYSBZ 2243-染色 （树链剖分）

#include <iostream>
#include <cstring>
#include <cstdio>

using namespace std;
const int N = 100005;
int a[N];

struct Edge {
    int to, next;
} edge[N*2];
int head[N], cntE;

void addedge(int u, int v) {
    edge[cntE].to = v; edge[cntE].next = head[u]; head[u] = cntE++;
    edge[cntE].to = u; edge[cntE].next = head[v]; head[v] = cntE++;
}

int dep[N], fa[N], sz[N], son[N];
void dfs1(int u, int pre, int d) {
    dep[u] = d;
    sz[u] = 1;
    fa[u] = pre;
    for (int i = head[u]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if (v != pre) {
            dfs1(v, u, d+1);
            sz[u] += sz[v];
            if (son[u] == -1 || sz[v] > sz[son[u]]) son[u] = v;
        }
    }
}

int top[N], dfn[N], rk[N], idx;
void dfs2(int u, int tp) {
    top[u] = tp;
    dfn[u] = ++idx;
    rk[idx] = u;
    if (son[u] == -1) return ;
    dfs2(son[u], tp);
    for (int i = head[u]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if (v != fa[u] && v != son[u]) {
            dfs2(v, v);
        }
    }
}

int tr[N<<2], ltr[N<<2], rtr[N<<2]; // tr[i] means the number of color segment
int fg[N<<2];
#define lson o<<1
#define rson o<<1|1

void pushup(int o) {
    ltr[o] = ltr[lson];
    rtr[o] = rtr[rson];
    tr[o] = tr[lson] + tr[rson];
    if (rtr[lson] == ltr[rson]) tr[o]--;
}

void pushdown(int o) {
    if (fg[o]) {
        tr[lson] = tr[rson] = 1;
        fg[lson] = fg[rson] = fg[o];
        ltr[lson] = ltr[rson] = ltr[o];
        rtr[lson] = rtr[rson] = rtr[o];
        fg[o] = 0;
    }
}

void build(int o, int l, int r) {
    fg[o] = 0;
    if (l == r) {
        tr[o] = 1;
        ltr[o] = rtr[o] = a[rk[l]];
        return;
    }
    int mid = (l+r) >> 1;
    build(lson, l, mid);
    build(rson, mid+1, r);
    pushup(o);
}

void change(int o, int l, int r, int L, int R, int v) {
    if (l >= L && r <= R) {
        fg[o] = 1;
        tr[o] = 1;
        ltr[o] = rtr[o] = v;
        return ;
    }
    pushdown(o);
    int mid = (l+r) >> 1;
    if (mid >= L) change(lson, l, mid, L, R, v);
    if (mid < R) change(rson, mid+1, r, L, R, v);
    pushup(o);
}

void CHANGE(int x, int y, int n, int c) {
    while (top[x] != top[y]) {
        if (dep[top[x]] < dep[top[y]]) swap(x, y);
        change(1, 1, n, dfn[top[x]], dfn[x], c);
        x = fa[top[x]];
    }
    if (dep[x] > dep[y]) swap(x, y);
    change(1, 1, n, dfn[x], dfn[y], c);
}

int query(int o, int l, int r, int L, int R) {
    if (l >= L && r <= R) return tr[o];
    pushdown(o);
    int mid = (l+r) >> 1;
    if (mid < L) {
        return query(rson, mid+1, r, L, R);
    } else if (mid >= R) {
        return query(lson, l, mid, L, R);
    } else {
        int ans = query(lson, l, mid, L, R);
        ans += query(rson, mid+1, r, L, R);
        if (ltr[rson] == rtr[lson]) ans--;
        return ans;
    }
}

int qq(int o, int l, int r, int p) {
    if (l == r) return ltr[o];
    pushdown(o);
    int mid = (l+r) >> 1;
    if (mid >= p) return qq(lson, l, mid, p);
    return qq(rson, mid+1, r, p);
}

int QUERY(int x, int y, int n) {
    int ans = 0;
    while (top[x] != top[y]) {
        if (dep[top[x]] < dep[top[y]]) swap(x, y);
        ans += query(1, 1, n, dfn[top[x]], dfn[x]);
        if (qq(1, 1, n, dfn[top[x]]) == qq(1, 1, n, dfn[fa[top[x]]])) --ans;
        x = fa[top[x]];
    }
    if (dep[x] > dep[y]) swap(x, y);
    ans += query(1, 1, n, dfn[x], dfn[y]);
    return ans;
}

void init() {
    memset(head, -1, sizeof head);
    memset(son, -1, sizeof son);
    idx = cntE = 0;
}

int main()
{
    //freopen("in.txt", "r", stdin);
    int n, m;
    while (~scanf("%d%d", &n, &m)) {
        init();
        for (int i = 1; i <= n; ++i) scanf("%d", &a[i]);
        int u, v, c;
        for (int i = 1; i < n; ++i) {
            scanf("%d%d", &u, &v);
            addedge(u, v);
        }
        dfs1(1, 0, 0); dfs2(1, 1); build(1, 1, n);

        char op[10];
        while (m--) {
            scanf("%s", op);
            if (*op == ‘Q‘) {
                scanf("%d%d", &u, &v);
                printf("%d\n", QUERY(u, v, n));
            } else {
                scanf("%d%d%d", &u, &v, &c);
                CHANGE(u, v, n, c);
            }
        }
    }
    return 0;
}