牛客网NOIP提高组模拟赛第六场C-树solution

纪念一下目前写的最长的代码—5.66KB

本题就是巨型数据结构题，用树剖加线段树维护一下父亲边与儿子边的修改，换根分类讨论一下即可

直接上代码：

#include<cstdio>
#include<iostream>
namespace fast_IO {
const int IN_LEN=10000000,OUT_LEN=10000000;
char ibuf[IN_LEN],obuf[OUT_LEN],*ih=ibuf+IN_LEN,*oh=obuf;
char *lastin=ibuf+IN_LEN; const char *lastout=ibuf+OUT_LEN-1;
inline char getchar_() { if(ih==lastin)lastin=ibuf+fread(ibuf,1,IN_LEN,stdin),ih=ibuf; return (*ih++); }
inline void putchar_(const char x) { if(ih==lastout)fwrite(obuf,1,oh-obuf,stdout),oh=obuf; *oh++=x; }
inline void flush(){fwrite(obuf, 1, oh - obuf, stdout);} }
using namespace fast_IO;
#define getchar() getchar_()
#define putchar(x) putchar_((x))
inline int read(){int x=0;char cu=getchar();x=0;while(!isdigit(cu))cu=getchar(); while(isdigit(cu))x=x*10+cu-'0',cu=getchar();return x;}
template <typename T> void print(const T x){if(x>=10)print(x/10); putchar(x%10+'0'); }
using namespace std;
const int N=1e5+10;
int n,m,A,B,head[N],cnt,son[N],hson[N],depth[N],f[N],pre[N],las[N],cl,top[N],opt,x,y,col,colk,num;
struct edge{int nxt,to;}ed[N<<1];
struct data{int col,t;}t1,t2;
struct Segment_Tree{int sum[3];data tag,son;}seg[N<<2];
inline void addedge(int x,int y){
    ed[++cnt].to=y;ed[cnt].nxt=head[x];head[x]=cnt;
    ed[++cnt].to=x;ed[cnt].nxt=head[y];head[y]=cnt;}
inline void DFS1(int u,int fa,int de){
    depth[u]=de;f[u]=fa;son[u]=1;
    for(register int i=head[u];i;i=ed[i].nxt){
        int v=ed[i].to;if(v==fa)continue;
        DFS1(v,u,de+1);if(son[hson[u]]<son[v])hson[u]=v;son[u]+=son[v];}
}
inline void DFS2(int u,int heavy){
    pre[u]=++cl;top[u]=heavy;if(hson[u])DFS2(hson[u],heavy);
    for(register int i=head[u];i;i=ed[i].nxt){
        int v=ed[i].to;if(v==f[u]||v==hson[u])continue;DFS2(v,v);}
    las[u]=cl;
}
inline void build(int o,int L,int R){
    if(L==R){seg[o].tag=(data){0,0};seg[o].sum[0]=1;return;}
    int M=(L+R)>>1,lc=o<<1,rc=(o<<1)+1;build(lc,L,M);build(rc,M+1,R);
    seg[o].sum[0]=seg[lc].sum[0]+seg[rc].sum[0];seg[o].tag=seg[o].son=(data){-1,-1};
}
inline void pushdown(int o,int L,int R){
    if(seg[o].tag.col!=-1&&seg[o].tag.t!=-1){
        int M=(L+R)>>1,lc=o<<1,rc=(o<<1)+1;
        seg[lc].tag=seg[o].tag;seg[rc].tag=seg[o].tag;
        for(register int i=0;i<=2;i++)
            if(i==seg[o].tag.col)seg[lc].sum[i]=M-L+1;
            else seg[lc].sum[i]=0;
        for(register int i=0;i<=2;i++)
            if(i==seg[o].tag.col)seg[rc].sum[i]=R-M;
            else seg[rc].sum[i]=0;
        seg[o].tag=(data){-1,-1};
    }
    if(seg[o].son.col!=-1&&seg[o].son.t!=-1){
        int M=(L+R)>>1,lc=o<<1,rc=(o<<1)+1;
        seg[lc].son=seg[o].son;seg[rc].son=seg[o].son;
        seg[o].son=(data){-1,-1};
    }
}
inline int Query(int o,int L,int R,int l,int r,int v){
    if(l<=L&&R<=r)return seg[o].sum[v];int sum=0;
    pushdown(o,L,R);int M=(L+R)>>1,lc=o<<1,rc=(o<<1)+1;
    if(l<=M)sum+=Query(lc,L,M,l,r,v);if(r>M)sum+=Query(rc,M+1,R,l,r,v);
    return sum;}
inline data&Find(int o,int L,int R,int x,int op){
    if(L==R){if(op)return seg[o].tag;else return seg[o].son;}
    pushdown(o,L,R);int M=(L+R)>>1,lc=o<<1,rc=(o<<1)+1;
    if(x<=M)return Find(lc,L,M,x,op);else return Find(rc,M+1,R,x,op);
}
inline int query(int u,int v,int col){
    int ans=0;
    while(top[u]!=top[v]){
        if(depth[top[u]]<depth[top[v]])swap(u,v);
        if(pre[top[u]]+1<=pre[u])ans+=Query(1,1,n,pre[top[u]]+1,pre[u],col);
        data k1=Find(1,1,n,pre[top[u]],1),k2;
        if(pre[f[top[u]]])k2=Find(1,1,n,pre[f[top[u]]],0);else k2=(data){-1,-1};
        if(k2.t==-1||k1.t>=k2.t)ans+=(k1.col==col);else ans+=(k2.col==col);
        //k1.t==k2.t的时候算是k1后更新，取k1的颜色，因为链更新时将一条链上的son与tag一同更新，这样会导致son和tag的时间戳相同，然而实际上链上的点并不会受son的影响，它们的data依然是tag
        u=f[top[u]];
    }
    if(depth[u]<depth[v])swap(u,v);
    if(pre[v]+1<=pre[u])ans+=Query(1,1,n,pre[v]+1,pre[u],col);
    return ans;
}
inline void updata(int o,int L,int R,int l,int r,int op,data v){
    if(l<=L&&R<=r){
        if(op)seg[o].son=v;
        else{seg[o].tag=v;for(register int i=0;i<=2;i++)if(v.col==i)seg[o].sum[i]=R-L+1;else seg[o].sum[i]=0;}
        return;
    }
    pushdown(o,L,R);int M=(L+R)>>1,lc=o<<1,rc=(o<<1)+1;
    if(l<=M)updata(lc,L,M,l,r,op,v);if(r>M)updata(rc,M+1,R,l,r,op,v);
    for(register int i=0;i<=2;i++)seg[o].sum[i]=seg[lc].sum[i]+seg[rc].sum[i];
}
inline void update(int u,int v,int col1,int col2){
    num++;t1=(data){col1,num};t2=(data){col2,num};
    while(top[u]!=top[v]){
        if(depth[top[u]]<depth[top[v]])swap(u,v);
        updata(1,1,n,pre[top[u]],pre[u],0,t1);
        updata(1,1,n,pre[top[u]],pre[u],1,t2);
        if(hson[u])updata(1,1,n,pre[u]+1,pre[u]+1,0,t2);
        u=f[top[u]];
    }
    if(depth[u]<depth[v])swap(u,v);
    if(pre[v]+1<=pre[u])updata(1,1,n,pre[v]+1,pre[u],0,t1);
    updata(1,1,n,pre[v],pre[u],1,t2);
    if(hson[u])updata(1,1,n,pre[u]+1,pre[u]+1,0,t2);
    updata(1,1,n,pre[v],pre[v],0,t2);
}
inline int Jump(int u,int goal){
    int w;
    if(depth[u]<depth[goal])return -1;
    while(top[u]!=top[goal]){
        if(depth[top[u]]<depth[top[goal]])return -1;w=top[u];u=f[top[u]];}
    if(depth[u]<depth[goal])return -1;
    if(u==goal)return w;else return hson[goal];
}
int main(){
    n=read();
    for(register int i=1;i<n;i++)A=read(),B=read(),addedge(A,B);
    DFS1(1,0,1);DFS2(1,1);build(1,1,n);
    m=read();
    while(m--){
        opt=read();x=read();y=read();col=read();
        if(opt==2)colk=read();
        if(opt==1)printf("%d\n",query(x,y,col));
        if(opt==2)update(x,y,col,colk);
        if(opt==3){
            num++;t1=(data){col,num};
            if(x==y)updata(1,1,n,1,n,0,t1);
            else{
                int k=Jump(x,y);
                if(k==-1)updata(1,1,n,pre[y]+1,las[y],0,t1);
                else{
                    if(1<=pre[k]-1)updata(1,1,n,1,pre[k]-1,0,t1);
                    if(las[k]+1<=n)updata(1,1,n,las[k]+1,n,0,t1);}
            }
        }
    }
    flush();return 0;
}