HYSBZ 2243 染色 树链剖分 点上剖分

题目链接:http://bak2.vjudge.net/problem/38096/origin
题意:
两种操作:
1.修改u到v路径上的点的颜色为c
2.询问u到v路径上有多少段颜色

线段树上用区间合并来维护答案
主要是在询问时要考虑一下，维护上个段的最后节点颜色来维护颜色段数

#include <bits/stdc++.h>
#define sf scanf
#define pf printf

using namespace std;
const int maxn = 100000 + 5;
int n,m;
//边集
struct Edge{
    int v,pre,c;
}Es[maxn * 2];
int head[maxn],TOT_EDGE;
void INIT_EDGE(){memset(head,-1,sizeof head);TOT_EDGE = 0;}
void ADD_EDGE(int u,int v,int c){
    Es[TOT_EDGE].v = v;
    Es[TOT_EDGE].c = c;
    Es[TOT_EDGE].pre = head[u];
    head[u] = TOT_EDGE++;
}

//树链剖分
int son[maxn],fa[maxn],dep[maxn],size[maxn];
int top[maxn],tid[maxn],rank[maxn],SegSize;
int DFS(int rt){
    dep[rt] = dep[fa[rt]] + 1;
    son[rt] = 0;
    size[rt] = 1;
    for(int i = head[rt];~i;i = Es[i].pre){
        int v = Es[i].v;
        if(v != fa[rt]){
            fa[v] = rt;
            size[rt] += DFS(v);
            if(size[son[rt]] < size[v]) son[rt] = v;
        }
    }
    return size[rt];
}
void Split(int rt,int tp){
    top[rt] = tp;
    tid[rt] = ++SegSize;
    rank[tid[rt]] = rt;
    if(son[rt]){
        Split(son[rt],tp);
        for(int i = head[rt];~i;i = Es[i].pre){
            int v = Es[i].v;
            if(v != fa[rt] && v != son[rt]) Split(v,v);
        }
    }
}
void TreeLineSplit(){
    dep[0] = 0;
    size[0] = 0;
    fa[1] = 0;
    DFS(1);
    SegSize = 0;
    Split(1,1);
}

//线段树
#define lson rt << 1 , l , mid
#define rson rt << 1 | 1,mid + 1,r
int LEFT[maxn << 2],RIGHT[maxn << 2],CNT[maxn << 2],num[maxn];
int cover[maxn << 2];
void PushUp(int rt){
    LEFT[rt] = LEFT[rt << 1];
    RIGHT[rt] = RIGHT[rt << 1 | 1];
    CNT[rt] = CNT[rt << 1] + CNT[rt << 1 | 1] - (RIGHT[rt << 1] == LEFT[rt << 1 | 1]);
}
void PushDown(int rt){
    if(cover[rt] != -1){
        LEFT[rt << 1] = RIGHT[rt << 1] = cover[rt];
        LEFT[rt << 1 | 1] = RIGHT[rt << 1 | 1] = cover[rt];
        CNT[rt << 1] = CNT[rt << 1 | 1] = 1;
        cover[rt << 1] = cover[rt << 1 | 1] = cover[rt];
        cover[rt] = -1;
    }
}
void SegTreeBuild(int rt,int l,int r){
    cover[rt] = -1;
    if(l == r){
        LEFT[rt] = RIGHT[rt] = num[l];
        CNT[rt] = 1;
        return;
    }
    int mid = l + r >> 1;
    SegTreeBuild(lson);SegTreeBuild(rson);
    PushUp(rt);
}
int LEFT_COLOR,RIGHT_COLOR;
int Query(int rt,int l,int r,int L,int R){
    if(L <= l && R >= r){
        if(LEFT_COLOR == -1) LEFT_COLOR = LEFT[rt];
        RIGHT_COLOR = RIGHT[rt];
        return CNT[rt];
    }
    PushDown(rt);
    int mid = l + r >> 1,ret = 0,f = 0;
    if(L <= mid) ret += Query(lson,L,R),f++;
    if(R > mid) ret += Query(rson,L,R),f++;
    if(f == 2 && RIGHT[rt << 1] == LEFT[rt << 1 | 1]) ret--;
    PushUp(rt);
    return ret;
}
void Update(int rt,int l,int r,int L,int R,int c){
    if(L <= l && R >= r){
        cover[rt] = c;
        LEFT[rt] = RIGHT[rt] = c;
        CNT[rt] = 1;
        return;
    }
    PushDown(rt);
    int mid = l + r >> 1;
    if(L <= mid) Update(lson,L,R,c);
    if(R > mid) Update(rson,L,R,c);
    PushUp(rt);
}
int INIT_COLOR[maxn];
void GetNumAr(){
    for(int i = 1;i <= n;++i){
        num[tid[i]] = INIT_COLOR[i];
    }
}


//题目操作
int Ans_Query(int u,int v){
    int ret = 0;
    int preU = -1,preV = -1;
    while(top[u] != top[v]){
        LEFT_COLOR = -1;
        if(dep[top[u]] > dep[top[v]]){
            ret += Query(1,1,n,tid[top[u]],tid[u]);
            if(RIGHT_COLOR == preU){
                ret--;
            }
            preU = LEFT_COLOR;
            u = fa[top[u]];
        }
        else{
            ret += Query(1,1,n,tid[top[v]],tid[v]);
            if(RIGHT_COLOR == preV){
                ret--;
            }
            preV = LEFT_COLOR;
            v = fa[top[v]];
        }
    }
    LEFT_COLOR = -1;
    if(dep[u] < dep[v]){
        ret += Query(1,1,n,tid[u],tid[v]);
        if(RIGHT_COLOR == preV) ret--;
        if(LEFT_COLOR == preU) ret--;
    }
    else{
        ret += Query(1,1,n,tid[v],tid[u]);
        if(RIGHT_COLOR == preU) ret--;
        if(LEFT_COLOR == preV) ret--;

    }
    return ret;
}

void CHANGE(int u,int v,int c){
    while(top[u] != top[v]){
        if(dep[top[v]] < dep[top[u]]) swap(u,v);
        Update(1,1,n,tid[top[v]],tid[v],c);
        v = fa[top[v]];
    }
    if(dep[v] < dep[u]) swap(u,v);
    Update(1,1,n,tid[u],tid[v],c);
}
char op[10];
int main(){
//    freopen("rand.txt","r",stdin);
    while( ~sf("%d%d",&n,&m) ){
        INIT_EDGE();
        for(int i = 1;i <= n;++i) sf("%d",&INIT_COLOR[i]);
        for(int i = 1;i < n;++i){
            int u,v;sf("%d%d",&u,&v);
            ADD_EDGE(u,v,0);
            ADD_EDGE(v,u,0);
        }
        TreeLineSplit();
        GetNumAr();
        SegTreeBuild(1,1,n);
        LEFT_COLOR = -1;
        while(m--){
            int x,y,z;
            sf("%s",op);
            if(op[0] == 'Q'){
                sf("%d%d",&x,&y);
                pf("%d\n",Ans_Query(x,y));
            }
            else{
                sf("%d%d%d",&x,&y,&z);
                CHANGE(x,y,z);
            }
        }
    }
}