BZOJ 2243: [SDOI2011]染色

Description

一棵树求路径上颜色个数，支持修改。

Solution

树链剖分...

我差不多是个zz了...又忘下放标记...

Code

/**************************************************************
    Problem: 2243
    User: BeiYu
    Language: C++
    Result: Accepted
    Time:4972 ms
    Memory:39044 kb
****************************************************************/
 
#include <bits/stdc++.h>
using namespace std;
 
#define mpr make_pair
#define x first
#define y second
#define debug(a) cout<<(#a)<<"="<<a<<" "
#define lc (o<<1)
#define rc (o<<1|1)
#define mid ((l+r)>>1)
 
typedef long long LL;
typedef pair<int,int> pr;
typedef vector<int> Vi;
typedef vector<LL> Vl;
typedef vector<string> Vs;
const int N = 300005;
const int M = N<<2;
const int oo = 0x3fffffff;
const LL  OO = 1e18;
 
inline LL in(LL x=0,char ch=getchar(),int v=1) {
    while(ch>'9' || ch<'0') v=ch=='-'?-1:v,ch=getchar();
    while(ch>='0' && ch<='9') x=x*10+ch-'0',ch=getchar();
    return x*v;
}
/*end*/
struct Tr { int x,l,r; };
Tr operator + (const Tr &a,const Tr &b) {
    if(!a.x) return b;
    if(!b.x) return a;
    if(a.r==b.l) return (Tr) { a.x+b.x-1,a.l,b.r };
    else return (Tr) { a.x+b.x,a.l,b.r };
}
 
void out(const Tr a) { cout<<a.x<<" "<<a.l<<" "<<a.r<<endl; }
 
int n,m;
int p[N],rp[N],c[N];
 
struct Segment {
    Tr d[M];
    int tg[M];
     
    void Update(int o) { d[o]=d[lc]+d[rc]; }
    void Push(int o) {
        if(tg[o]) {
            if(lc) tg[lc]=tg[o],d[lc]=(Tr) { 1,tg[o],tg[o] };
            if(rc) tg[rc]=tg[o],d[rc]=(Tr) { 1,tg[o],tg[o] };
            tg[o]=0;
        }
    }
    void Build(int o,int l,int r) {
        if(l==r) { d[o]=(Tr) { 1,c[rp[l]],c[rp[l]] };return; }
        Build(lc,l,mid),Build(rc,mid+1,r);
        Update(o);
    }
    void Modify(int o,int l,int r,int L,int R,int v) {
        if(L<=l && r<=R) { d[o]=(Tr) { 1,v,v },tg[o]=v;return; }
        Push(o);
        if(L<=mid) Modify(lc,l,mid,L,R,v);
        if(R>mid) Modify(rc,mid+1,r,L,R,v);
        Update(o);
    }
    Tr Query(int o,int l,int r,int L,int R) {
        Push(o);
        if((L<=l && r<=R)||d[o].x==1) return d[o];
        if(L<=mid && R>mid) return Query(lc,l,mid,L,R)+Query(rc,mid+1,r,L,R);
        if(L<=mid) return Query(lc,l,mid,L,R);
        else return Query(rc,mid+1,r,L,R);
    }
}py;
 
namespace Tree {
    int cnt;
    int d[N],top[N],sn[N],sz[N],f[N];
    vector<int> g[N];
     
    void AddEdge(int u,int v) { g[u].push_back(v),g[v].push_back(u); }
    void DFS1(int u,int fa) {
        d[u]=d[fa]+1,sz[u]=1,sn[u]=0;
        for(int i=0,v;i<(int)g[u].size();i++) if((v=g[u][i])!=fa) {
            DFS1(v,u),sz[u]+=sz[v],f[v]=u;
            if(!sn[u] || sz[sn[u]]<sz[v]) sn[u]=v;
        }
    }
    void DFS2(int u,int fa,int tp) {
        top[u]=tp,p[u]=++cnt,rp[cnt]=u;
        if(sn[u]) DFS2(sn[u],u,tp);
        for(int i=0,v;i<(int)g[u].size();i++) if((v=g[u][i])!=fa && v!=sn[u])
            DFS2(v,u,v);
    }
    void init() {
        DFS1(1,1);
        DFS2(1,1,1);
        py.Build(1,1,n);
    }
    void Modify(int u,int v,int c) {
        int f1=top[u],f2=top[v];
        for(;f1^f2;) {
            if(d[f1]<d[f2]) swap(f1,f2),swap(u,v);
            py.Modify(1,1,n,p[f1],p[u],c);
            u=f[f1],f1=top[u];
        }if(d[u]>d[v]) swap(u,v);
        py.Modify(1,1,n,p[u],p[v],c);
    }
    int Query(int u,int v) {
        int f1=top[u],f2=top[v],r=0;
        Tr r1,r2;r1=r2=(Tr){ 0,-1,-1 };
        for(;f1^f2;) {
            if(d[f1]<d[f2]) swap(f1,f2),swap(u,v),r^=1;
            if(r) r2=py.Query(1,1,n,p[f1],p[u])+r2;
            else r1=py.Query(1,1,n,p[f1],p[u])+r1;
            u=f[f1],f1=top[u];
        }
        if(d[u]>d[v]) swap(u,v),r^=1;
        if(r) r1=py.Query(1,1,n,p[u],p[v])+r1;
        else r2=py.Query(1,1,n,p[u],p[v])+r2;
        swap(r1.l,r1.r);
        return (r1+r2).x;
    }
}
 
int main() {
    n=in(),m=in();
    for(int i=1;i<=n;i++) c[i]=in()+1;
    for(int i=1;i<n;i++) {
        int x=in(),y=in();
        Tree::AddEdge(x,y);
    }
    Tree::init();
    for(int i=1;i<=m;i++) {
        char opt[15];
        scanf("%s",opt);
        if(opt[0]=='C') {
            int u=in(),v=in(),c=in()+1;
            Tree::Modify(u,v,c);
        } else {
            int u=in(),v=in();
            printf("%d
",Tree::Query(u,v));
        }
    }return 0;
}