bzoj2243 樹刨+線段樹區間合併

題意：區間修改樹上兩點間的顏色種類，區間查詢樹上兩點間的線段數。
思路：樹刨+線段樹，區間合併是第一次寫。思路爲線段樹維護左端點和右端點的顏色，合併時相同就–。具體實現不太好寫。
代碼：

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define forn(i,n) for(int i=0;i<n;i++)
#define for1(i,n) for(int i=1;i<=n;i++)
#define IO ios::sync_with_stdio(false);cin.tie(0)
const int maxn = 1e5+5;

int a[maxn],top[maxn],sz[maxn],son[maxn],deep[maxn],p[maxn],par[maxn],fp[maxn],id,now;
vector<int>e[maxn];
class segment_tree{
    public:
    #define nd node[now]
    #define ndl node[now<<1]
    #define ndr node[now<<1|1]
    struct segment_node{
        int l,r,v,flag,cl,cr;
        void update(int x){
            flag = cl = cr = x,v = 1;
        }
    }node[maxn<<2];
    void pushup(int now){
		nd.cl = ndl.cl,nd.cr = ndr.cr; 
        nd.v = ndl.v+ndr.v;
        if(ndl.cr==ndr.cl) nd.v--;
    }
    void pushdown(int now){
        if(nd.flag!=-1){
            ndl.update(nd.flag);
            ndr.update(nd.flag);
            nd.flag = -1;
        }
    }
    void maketree(int l,int r,int now = 1){
        nd = {l,r,-1,-1,-1,-1};
        if(l==r){
            nd.v = 1,nd.cl = nd.cr = a[fp[l]];
            return;
        }
        maketree(l,l+r>>1,now<<1);
        maketree((l+r>>1)+1,r,now<<1|1);
        pushup(now);
    }
    void update(int l,int r,int v,int now = 1){
        if(l<=nd.l&&r>=nd.r){
            nd.update(v);
            return;
        }
        pushdown(now);
        if(l<=ndl.r) update(l,r,v,now<<1);
        if(r>=ndr.l) update(l,r,v,now<<1|1);
        pushup(now);
    }
    int query(int l,int r,int now = 1){
        //cerr<<l<< ' '<<r<<' '<<now<<'$'<<nd.l<<' '<<nd.r<<'\n';
        if(l<=nd.l&&r>=nd.r) return nd.v;
        int res = 0;
        pushdown(now);
        if(l<=ndl.r&&r>=ndr.l){
            res+=query(l,r,now<<1);
            res+=query(l,r,now<<1|1);
            if(ndl.cr==ndr.cl) res--;
        }
        else if(l<=ndl.r) res+=query(l,r,now<<1);
        else if(r>=ndr.l) res+=query(l,r,now<<1|1);
        pushup(now);
        return res;
    }
    int getcr(int pos,int now = 1){
        if(nd.l==nd.r) return nd.cr;
        int res = 0;
        pushdown(now);
        if(pos<=ndl.r) res = getcr(pos,now<<1);
        else res = getcr(pos,now<<1|1);
        pushup(now);
        return res;
    }
}tree;
void dfs(int u,int pre,int d){
    deep[u] = d,par[u] = pre,sz[u] = 1;
    forn(i,e[u].size()){
        int v = e[u][i];
        if(v==pre) continue;
        dfs(v,u,d+1);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u] = v;
    }
}
void getpos(int u,int gg){
    top[u] = gg,p[u] = ++id,fp[p[u]] = u;
    if(son[u]) getpos(son[u],gg);
    forn(i,e[u].size()){
        int v = e[u][i];
        if(v==par[u]||v==son[u]) continue;
        getpos(v,v);
    }
}
bool check(int x,int y){
    return tree.getcr(p[x])==tree.getcr(p[y]);
}
int query(int x,int y){
    int fx = top[x],fy = top[y];
    int res = 0;
    while(fx!=fy){
        if(deep[fx]<deep[fy]) swap(fx,fy),swap(x,y);
        res+=tree.query(p[fx],p[x]);
        if(check(par[fx],fx)) res--;
        x = par[fx], fx = top[x];
    }
    if(deep[x]>deep[y]) swap(x,y);
    res+=tree.query(p[x],p[y]);
    return res;
}

void change(int x,int y,int v){
    int fx = top[x],fy = top[y];
    while(fx!=fy){
        if(deep[fx]<deep[fy]) swap(fx,fy),swap(x,y);
        tree.update(p[top[x]],p[x],v);
        x = par[fx], fx = top[x];
    }
    if(deep[x]>deep[y]) swap(x,y);
    tree.update(p[x],p[y],v);
}

int main(){
    IO;
    int n,m;cin>>n>>m;
    for1(i,n) cin>>a[i];
    forn(i,n-1){
        int x,y;cin>>x>>y;
        e[x].push_back(y);
        e[y].push_back(x);
    }
    dfs(1,0,0);    
    getpos(1,1);
    tree.maketree(1,n);
    while(m--){
        char c;cin>>c;
        if(c=='Q'){
            int x,y;cin>>x>>y;
            cout<<query(x,y)<<'\n';
        }else{
            int x,y,z;cin>>x>>y>>z;
            change(x,y,z);
        }   
    }
    return 0;
}