【BZOJ 2243】[SDOI2011]染色

题目来源：BZOJ 2243

思路：

在处理这样的题的时候，为了看是否可以用树链剖分，优先考虑如果操作是在序列上的时候应该怎么做。
观察两个操作：
1、将节点a到节点b路径上所有点都染成颜色c；
2、询问节点a到节点b路径上的颜色段数量。
对于第一个操作，可以打颜色的标记，用线段树维护区间。
对于第二个操作，可以记录一个区间内的颜色个数，以及区间内最左的颜色和最右的颜色。
考虑区间信息的合并，如果左区间的右端颜色与右区间的左端颜色不同，新区间的颜色段数直接相加，如果颜色相同，新区间的颜色个数是相加后-1。
知道了这些就可以考虑树链剖分了。
考虑查询的方法：对于连续重链上的一段区间，合并查询到的每一个小区间的信息，返回个数，左端点，右端点；不在同一个条重链上的边，将这个区间的右端点与上一个区间的左端点合并，变成一个新的区间信息；对于最后一条公共链来说，讨论左右端点的关系，进行合并。
注意：0也是一种颜色。

代码：

#include <cstdio>
#include <iostream>
#include <algorithm>
#define mid ((l+r)>>1)
#define lch (now<<1)
#define rch ((now<<1)+1) 
const int maxn = 100010;
const int inf = 2e9;
struct node{
    int num, ll, rr;
} t_node[maxn<<2], tmp, tp;
int n, m, ct;
int head[maxn], next[maxn<<1], to[maxn<<1], cnt;
int size[maxn], fa[maxn], son[maxn], dep[maxn];
int pos[maxn], top[maxn], v[maxn], val[maxn];
int lazy[maxn<<2];
void add(int a, int b){
    next[++cnt] = head[a];
    to[head[a] = cnt] = b;
}
void dfs1(int x){
    size[x] = 1;
    for(int i = head[x]; i; i = next[i]){
        if(to[i] == fa[x]) continue;
        dep[to[i]] = dep[x] + 1, fa[to[i]] = x;
        dfs1(to[i]), size[x] += size[to[i]];
        if(size[son[x]] < size[to[i]]) son[x] = to[i]; 
    }
}
void dfs2(int x){
    pos[x] = ++ct;
    if(son[x]) top[son[x]] = top[x], dfs2(son[x]);
    for(int i = head[x]; i; i = next[i]){
        if(to[i] == fa[x] || to[i] == son[x]) continue;
        top[to[i]] = to[i], dfs2(to[i]);
    }
}
inline node merge(node a, node b){
    if(a.num == 0) return b;
    node res;
    res.num = a.num + b.num;
    res.ll = a.ll, res.rr = b.rr;
    if(a.rr == b.ll) res.num -= 1;
    return res;
}
// 颜色区间有0 
void down(int now, int l, int r){
    if(lazy[now] == -1) return;
    t_node[lch].ll = t_node[lch].rr = t_node[rch].ll = t_node[rch].rr = lazy[now];
    lazy[lch] = lazy[rch] = lazy[now];
    t_node[lch].num = t_node[rch].num = 1;
    lazy[now] = -1;
} 
void build(int now, int l, int r){
    lazy[now] = -1; 
    if(l == r){t_node[now].num = 1, t_node[now].ll = v[l], t_node[now].rr = v[l]; return;}
    build(lch, l, mid);
    build(rch, mid+1, r);
    t_node[now] = merge(t_node[lch], t_node[rch]);
}
void modify(int now, int l, int r, int pos1, int pos2, int k){
    if(l == pos1 && r == pos2){
        t_node[now].num = 1, lazy[now] = k;
        t_node[now].ll = k, t_node[now].rr = k;
        return;
    }down(now, l, r);
    if(pos2 <= mid) modify(lch, l, mid, pos1, pos2, k);
    else if(pos1 >= mid+1) modify(rch, mid+1, r, pos1, pos2, k);
    else modify(lch, l, mid, pos1, mid, k), modify(rch, mid+1, r, mid+1, pos2, k);
    t_node[now] = merge(t_node[lch], t_node[rch]); 
}
node que_num(int now, int l, int r, int pos1, int pos2){
    if(l == pos1 && r == pos2){return t_node[now];}
    down(now, l, r);
    if(pos2 <= mid) return que_num(lch, l, mid, pos1, pos2);
    else if(pos1 >= mid+1) return que_num(rch, mid+1, r, pos1, pos2);
    else return merge(que_num(lch, l, mid, pos1, mid), que_num(rch, mid+1, r, mid+1, pos2));
}
int que_num(int x, int y){
    node res1 = tp, res2 = tp;
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]){
            std::swap(x, y);
            std::swap(res1, res2);
        } 
        res1 = merge(que_num(1, 1, n, pos[top[x]], pos[x]), res1);
        x = fa[top[x]];
    }
    if(pos[x] > pos[y]) std::swap(x, y);
    else std::swap(res1, res2);
    res1 = merge(que_num(1, 1, n, pos[x], pos[y]), res1);
    if(res1.ll == res2.ll) return res1.num + res2.num - 1;
    else return res1.num + res2.num;
}
void modify(int x, int y, int k){
    while(top[x] != top[y]){
        if(dep[top[x]] < dep[top[y]]) std::swap(x, y);
        modify(1, 1, n, pos[top[x]], pos[x], k);
        x = fa[top[x]];
    }
    if(pos[x] > pos[y]) std::swap(x, y);
    modify(1, 1, n, pos[x], pos[y], k);
}
int main(){
    tp.num = 0, tp.ll = tp.rr = -1;
    scanf("%d%d", &n, &m);
    for(int i = 1; i <= n; i ++) scanf("%d", &val[i]);
    for(int i = 1, a, b; i < n; i ++){
        scanf("%d%d", &a, &b);
        add(a, b), add(b, a);
    }
    dep[1] = 1, dfs1(1);
    top[1] = 1, dfs2(1);
    for(int i = 1; i <= n; i ++) v[pos[i]] = val[i];
    build(1, 1, n);
    for(int i = 1, a, b, c; i <= m; i ++){
        char ch[5];
        scanf("%s", ch);
        if(ch[0] == 'Q'){
            scanf("%d%d", &a, &b);
            printf("%d\n", que_num(a, b));
        }else{
            scanf("%d%d%d", &a, &b, &c);
            modify(a, b, c);
        }
    }
    return 0;
}