【BZOJ】4771: 七彩树 -线段树合并&主席树&set

传送门：bzoj4771

---

题解

为满足深度和子树两个限制，预处理出 dfs d f s 序和每层的点,然后按深度递增， dfs d f s 序为下标建立主席树，每次询问查询 (rt[depx+d],in[x],ot[x]) ( r t [ d e p x + d ] , i n [ x ] , o t [ x ] ) ， [in[x],ot[x]] [ i n [ x ] , o t [ x ] ] 表示 dfs d f s 序中 x x 的子树区间。
为保证每个出现的颜色贡献为 1 1 ，每个颜色建一个按 dfs d f s 序升序排列的 set s e t ， set s e t 中每个点贡献 +1 + 1 ，每两个邻点的 LCA L C A 贡献 −1 − 1 ，每次增加深度更新 set s e t 时， lower l o w e r _ bound b o u n d 找到位置，把前驱和后继的 LCA L C A 贡献 +1 + 1 ，再分别把该点和前驱后继的 LCA L C A 贡献 +1 + 1 即可。

---

代码

#include
#define gc getchar()
#define si isdigit(c)
#define mid (((l)+(r))>>1)
using namespace std;
const int N=1e5+10,M=1e7+10;

int T,n,m,ans,d[N],col[N],in[N],ot[N],dfn,cnt;
int rt[N],ls[M],rs[M],ss[M],mxdep,F[N][19],bin[20];

struct G{
   int tot,head[N],to[N],nxt[N];
   inline void clr()
   {memset(head,0,sizeof(head));tot=0;}
   inline void lk(int u,int v)
   {to[++tot]=v;nxt[tot]=head[u];head[u]=tot;}  
}A,B;

struct cmp{
    bool operator()(const int x,const int y){
       return in[x]set<int,cmp>S[N];

char c;
inline void rd(int &x)
{
    c=gc;x=0;
    for(;!si;c=gc);
    for(;si;c=gc) x=x*10+(c^48);
}

inline void dfs(int x)
{
    in[x]=++dfn;B.lk(d[x],x);
    int i;
    for(i=1;bin[i]<=d[x];++i)
     F[x][i]=F[F[x][i-1]][i-1];
    for(;i<19;++i) F[x][i]=0;
    for(i=A.head[x];i;i=A.nxt[i]){
       d[A.to[i]]=d[x]+1;F[A.to[i]][0]=x;dfs(A.to[i]);  
    }
    ot[x]=dfn;
    if(in[x]==ot[x]) mxdep=max(d[x],mxdep);
}

inline void pushup(int k)
{ss[k]=ss[ls[k]]+ss[rs[k]];}

inline void ins(int &k,int l,int r,int pos,int val)
{
    if(!k){k=++cnt;ls[k]=rs[k]=ss[k]=0;}
    if(l==r) {ss[k]+=val;return;}
    if(pos<=mid) ins(ls[k],l,mid,pos,val);
    else ins(rs[k],mid+1,r,pos,val);
    pushup(k);
}

inline int merge(int x,int y,int l,int r)
{
    if(!x || !y) return x|y;
    if(l==r) {ss[x]+=ss[y];return x;}
    ls[x]=merge(ls[x],ls[y],l,mid);
    rs[x]=merge(rs[x],rs[y],mid+1,r);
    pushup(x);
    return x;
}

inline int LCA(int x,int y)
{
    if(d[x]if(x==y) return x;
    int t=d[x]-d[y],i;
    for(i=0;bin[i]<=t;++i)
     if((t&bin[i])) x=F[x][i];
    for(i=18;i>=0 && x!=y;--i)
     if(F[x][i]!=F[y][i])
      x=F[x][i],y=F[y][i];
    return x==y?x:F[x][0];
}

inline int query(int k,int l,int r,int L,int R)
{
    if(L<=l && r<=R) return ss[k];
    int re=0;
    if(L<=mid) re=query(ls[k],l,mid,L,R);
    if(R>mid) re+=query(rs[k],mid+1,r,L,R);
    return re;
}

int main(){
    int i,j,x,D,nxt,pre;
    bin[0]=1;
    for(i=1;i<19;++i) bin[i]=bin[i-1]<<1;
    rd(T);
    while(T--){
        cnt=ans=dfn=0;A.clr();B.clr();
        rd(n);rd(m);
        for(i=1;i<=n;++i) rd(col[i]),S[i].clear();
        for(i=2;i<=n;++i){rd(x);A.lk(x,i);}
        d[1]=1;mxdep=0;dfs(1);
        rt[0]=0;
        for(D=1;D<=mxdep;++D){
            rt[D]=0;
            for(i=B.head[D];i;i=B.nxt[i]){
                j=B.to[i];
                ins(rt[D],1,n,in[j],1);
                set<int>::iterator ori,orz;
                ori=orz=S[col[j]].lower_bound(j);
                pre= ori==S[col[j]].begin()?-1:*(--ori);
                nxt= orz==S[col[j]].end()?-1:*orz;
                if(pre!=-1) ins(rt[D],1,n,in[LCA(pre,j)],-1);
                if(nxt!=-1) ins(rt[D],1,n,in[LCA(j,nxt)],-1);
                if(pre!=-1 && nxt!=-1) ins(rt[D],1,n,in[LCA(pre,nxt)],1);
                S[col[j]].insert(j);
            }
            rt[D]=merge(rt[D],rt[D-1],1,n);
        }
        while(m--){
            rd(x);rd(D);
            x^=ans;D^=ans;
            D=min(mxdep,d[x]+D);
            printf("%d\n",(ans=query(rt[D],1,n,in[x],ot[x])));
        }
    }
}