BZOJ 2243 [SDOI2011] 染色(树剖 + 线段树)

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

Input
第一行包含2个整数$n$和$m$，分别表示节点数和操作数；
第二行包含$n$个正整数表示$n$个节点的初始颜色
下面 行每行包含两个整数$x$和$y$，表示$x$和$y$之间有一条无向边。
下面 行每行描述一个操作：
“C a b c”表示这是一个染色操作，把节点$a$到节点$b$路径上所有点（包括$a$和$b$）都染成颜色$c$；
“Q a b”表示这是一个询问操作，询问节点$a$到节点$b$（包括$a$和$b$）路径上的颜色段数量。

Output
对于每个询问操作，输出一行答案。

Sample Input
6 5
2 2 1 2 1 1
1 2
1 3
2 4
2 5
2 6
Q 3 5
C 2 1 1
Q 3 5
C 5 1 2
Q 3 5
Sample Output
3
1
2

Hint
数$N<=10^5$，操作数$M<=10^5$，所有的颜色C为整数且在$[0,10^9]$之间。

Solution
看题意很明显树剖+线段树

线段树部分：
线段树除了维护区间内的颜色段数，还需要储存区间左端点的颜色和右端点的颜色，这样才能正确合并区间贡献
自然还需要lazy tag 来做区间修改

树剖部分：
一条重链内的答案可以通过线段树内的处理计算，但重链和重链之间因为序号不连续，所以无法通过线段树来计算
容易发现，重链转移时的 fa[top[u]] 和 top[u] 就是两次区间查询时的相邻端点，从而实现了两种情况的区间贡献合并

Code

#include 
#define ls rt << 1
#define rs rt << 1 | 1
using namespace std;

typedef long long ll;
const int MX = 2e5 + 7;

int n;
int ecnt,head[MX];
int w[MX],wt[MX];

struct Edge{
     
    int v,next;
}e[MX << 1];

void add(int u,int v){
     
    e[++ecnt].v = v;e[ecnt].next = head[u];head[u] = ecnt;
    e[++ecnt].v = u;e[ecnt].next = head[v];head[v] = ecnt;
}

struct Segtree{
     
    int l,r;
    int sum;
    int lc,rc;//区间左右端点颜色
    int lazy;
}t[MX << 2];

void push_up(int rt){
     
    t[rt].sum = t[ls].sum + t[rs].sum - (t[ls].rc == t[rs].lc);
    t[rt].lc = t[ls].lc, t[rt].rc = t[rs].rc;
}

void build(int rt,int l,int r){
     
    t[rt].l = l, t[rt].r = r;
    t[rt].lazy = -1;
    if(l == r){
     
        t[rt].sum = 1;
        t[rt].lc = t[rt].rc = wt[l];
        return ;
    }
    int mid = (l + r) >> 1;
    build(ls,l,mid);build(rs,mid+1,r);
    push_up(rt);
}

void push_down(int rt){
     
    if(t[rt].lazy != -1){
     
        int col = t[rt].lazy;
        t[ls].sum = 1; t[ls].lazy = col; t[ls].lc = t[ls].rc = col;
        t[rs].sum = 1; t[rs].lazy = col; t[rs].lc = t[rs].rc = col;
        t[rt].lazy = -1;
    }
}

void update(int rt,int L,int R,int k){
     
    int l = t[rt].l, r = t[rt].r;
    if(L <= l && r <= R){
     
        t[rt].sum = 1;
        t[rt].lc = t[rt].rc = k;
        t[rt].lazy = k;
        return ;
    }
    push_down(rt);
    int mid = (l + r) >> 1;
    if(L <= mid) update(ls,L,R,k);
    if(R  > mid) update(rs,L,R,k);
    push_up(rt);
}

int nowl;//跨重链合并区间贡献
int query(int rt,int L,int R){
     
    int l = t[rt].l, r = t[rt].r;
    if(L <= l && r <= R){
     
        if(L == l) nowl = t[rt].lc;
        return t[rt].sum;
    }
    push_down(rt);
    int mid = (l + r) >> 1, res = 0;
    if(L <= mid) res += query(ls,L,R);
    if(R  > mid) res += query(rs,L,R);
    if(L <= mid && R > mid){
     
        if(t[ls].rc == t[rs].lc) res--;
    }
    return res;
}

int siz[MX],son[MX],top[MX],dep[MX],fa[MX],id[MX],cnt;

void dfs1(int u,int f,int deep){
     
    siz[u] = 1, fa[u] = f,dep[u] = deep;
    int maxson = -1;
    for(int i = head[u];i;i = e[i].next){
     
        int v = e[i].v;
        if(v == f) continue;
        dfs1(v,u,deep+1);
        siz[u] += siz[v];
        if(siz[v] > maxson) maxson = siz[v], son[u] = v;
    }
}

void dfs2(int u,int topf){
     
    id[u] = ++cnt, wt[cnt] = w[u], top[u] = topf;
    if(!son[u]) return ;
    dfs2(son[u],topf);
    for(int i = head[u];i;i = e[i].next){
     
        int v = e[i].v;
        if(v == fa[u] || v == son[u]) continue;
        dfs2(v,v);
    }
}

void updRange(int u,int v,int k){
     
    while(top[u] != top[v]){
     
        if(dep[top[u]] < dep[top[v]]) swap(u,v);
        update(1,id[top[u]],id[u],k);
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u,v);
    update(1,id[u],id[v],k);
}

int qRange(int u,int v){
     
    int res = 0;
    while(top[u] != top[v]){
     
        if(dep[top[u]] < dep[top[v]]) swap(u,v);
        res += query(1,id[top[u]],id[u]);
        //跨重链合并区间贡献
        int now = nowl;
        int tmp = query(1,id[fa[top[u]]],id[fa[top[u]]]);
        int nx = nowl;
        if(nx == now) res--;
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u,v);
    res += query(1,id[u],id[v]);
    return res;
}

int main(){
     
    int q;
    scanf("%d%d",&n,&q);
    for(int i = 1;i <= n;++i) scanf("%d",&w[i]);
    for(int i = 1;i < n;++i) {
     
        int u,v;scanf("%d%d",&u,&v);add(u,v);
    }
    dfs1(1,-1,1);dfs2(1,1);
    build(1,1,n);
    while(q--){
     
        char op[4];scanf("%s",op);
        if(op[0] == 'C'){
     
            int u,v,k;scanf("%d%d%d",&u,&v,&k);
            updRange(u,v,k);
        } else if(op[0] == 'Q'){
     
            int u,v;scanf("%d%d",&u,&v);
            int res = qRange(u,v); printf("%d\n", res);
        }
    }
    return 0;
}