SDOI 2011 染色 树链剖分

题意：

给定一棵有n个节点的无根树和m个操作，操作有2类：

1）将节点a到节点b路径上所有点都染成颜色c；
2）询问节点a到节点n路径上的颜色段数量（连续相同颜色被认为是同一段），如“112221”由3段组成：“11”、“222”和“1”。

n, m <= 100000

思路：

树的形态没有改变，用树链剖分维护即可。

在线段树上维护区间内的颜色前缀，后缀以及颜色段数构成一个三元组（pre，suf，tot），则可很方便地合并信息：

(pre1, suf1, tot1) + (pre2, suf2, tot2) = (pre1, suf2, tot1 + tot2 - (suf1 == pre2))

注意一开始输入的颜色对应的点在线段树中的编号与原编号不同，查询操作合并信息要特别注意不要加反了。

又复习了一遍树剖，感觉细节比较多，写的时候要想清楚。。。

#include <cstdio>
#include <vector>

#define For(i,j,k) for(int i = j;i <= (int)k;i++)

#define lc (h << 1), L, M
#define rc (h << 1 | 1), M+1, R
#define root 1, 1, n
#define M ((L + R) >> 1)

using namespace std;

int n, m;

const int N = 100010;

int fa[N], id[N], top[N], co[N], c[N], sz[N], son[N], dep[N], e;
vector<int> G[N];

void DFS1(int h){
    sz[h] = 1, dep[h] = dep[fa[h]] + 1;
    For(i,0,G[h].size() - 1){
        int v = G[h][i];
        if(v == fa[h]) continue;
        fa[v] = h;
        DFS1(v);
        sz[h] += sz[v];
        if(sz[v] > sz[son[h]]) son[h] = v;
    }
}

void DFS2(int h){
    id[h] = ++e;
    c[e] = co[h];
    top[h] = son[fa[h]] == h ? top[fa[h]] : h;
    if(son[h]) DFS2(son[h]);
    For(i,0,G[h].size() - 1){
        int v = G[h][i];
        if(v == fa[h] || v == son[h]) continue;
        DFS2(v);
    }
}

struct info{
    int pre, suf, tot;

    info rev(){
        return (info){suf, pre, tot};
    }

    info operator + (const info& A) const{
        return (info){pre, A.suf, tot + A.tot - (suf == A.pre)};
    }
};

struct SegmentTree{
    info A[N<<2];
    int col[N<<2];

    void pushup(int h){
        A[h] = A[h << 1] + A[h << 1 | 1];
    }

    void pushdown(int h, int L, int R){
        if(col[h]){
            Modify(lc, L, M, col[h]), Modify(rc, M+1, R, col[h]);
            col[h] = 0;
        }
    }

    void Create(int h, int L, int R){
        if(L == R) A[h] = (info){c[L], c[L], 1}, col[h] = c[L];
        else{
            Create(lc), Create(rc);
            pushup(h);
        }
    }

    void Modify(int h, int L, int R, int l, int r, int v){
        if(l <= L && R <= r)
            A[h] = (info){v, v, 1}, col[h] = v;
        else{
            pushdown(h, L, R);
            if(l <= M) Modify(lc, l, r, v);
            if(r > M) Modify(rc, l, r, v);
            pushup(h);
        }
    }

    info Query(int h, int L, int R, int l, int r){
        if(l <= L && R <= r) return A[h];
        else{
            pushdown(h, L, R);
            if(r <= M) return Query(lc, l, r);
            else if(l > M) return Query(rc, l, r);
            else return Query(lc, l, r) + Query(rc, l, r);
        }
    }

}T;

void change(int u, int v, int val){
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u, v);
        T.Modify(1, 1, n, id[top[u]], id[u], val);
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u, v);
    T.Modify(root, id[u], id[v], val);
}

int query(int u, int v){
    info l, r;
    l = r = (info){0, 0, 0};
    int x = u, y = v;
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]) y = fa[top[y]];
        else x = fa[top[x]];
    }
    while(u != x) l = T.Query(root, id[top[u]], id[u]) + l, u = fa[top[u]];
    while(v != y) r = T.Query(root, id[top[v]], id[v]) + r, v = fa[top[v]];
    if(dep[u] > dep[v]) return (l.rev() + T.Query(1, 1, n, id[v], id[u]).rev() + r).tot;
    else return (l.rev() + T.Query(root, id[u], id[v]) + r).tot;
}

int main(){
    scanf("%d%d", &n, &m);
    For(i,1,n) scanf("%d", &co[i]);
    For(i,2,n){
        int x, y;
        scanf("%d%d", &x, &y);
        G[x].push_back(y), G[y].push_back(x);
    }
    DFS1(1), DFS2(1);
    T.Create(1, 1, n);
    while(m--){
        char Q[3];
        int u, v, x;
        scanf("%s%d%d", Q, &u, &v);
        if(Q[0] == 'C') scanf("%d", &x), change(u, v, x);
        else printf("%d\n", query(u, v));
    }
    return 0;
}