SP11985 GOT - Gao on a tree

\(\texttt{Description}\)

SPOJ 题目链接

洛谷题目链接

• 有一棵 \(n\) 个节点的树，第 \(i\) 个节点有颜色 \(a_i\)，\(q\) 次询问，每次询问 \(u,v\) 节点之间（含）是否存在颜色为 \(w\) 的节点。

• \(n\le 10^5\)，\(q\le 2\times 10^5\)，\(0\le a_i,w\le n\)。

\(\texttt{Solution}\)

和洛谷 P3313 一样，对于这类维护颜色的问题，可以考虑使用动态开点线段树来解决。

首先进行树链剖分把树上路径问题转化成序列问题，然后对每种颜色建立线段树。由于是存在性的问题，可以考虑建树时将每个点权值赋成 \(1\)，维护区间或和。

对于每次询问，在颜色 \(c\) 对应的线段树上跳链查询即可。

时间复杂度为 \(\mathcal{O}(q\log^2n)\)，空间复杂度为 \(\mathcal{O}(n\log n)\)，可以接受。

\(\texttt{Code}\)

代码实现时一些小细节：

• 多测清空。

• 我使用指针实现动态开点线段树，注意判左右儿子是不是 NULL。以及修改的时候参数传的是地址，否则修改的不是线段树本身，而是形参创造的一个副本（个人理解，轻喷）。

• 查询或合并左右儿子的时候如果遇到了 \(1\) 就直接返回，因为已经说明存在性了。

评测记录

点击查看代码

#pragma GCC optimize("Ofast")
#include<bits/stdc++.h>
#define N 100005
#define ls x->lson
#define rs x->rson
using namespace std;
int n,q,a[N],h[N],t[N],d[N],f[N],sz[N],dfn[N];
struct node{
    bool v;
    node*lson,*rson;
}*rt[N];
vector<int>g[N];
void dfs1(int u,int fa){
    sz[u]=1;
    for(int v:g[u]){
        if(v^fa){
            d[v]=d[u]+1;
            dfs1(v,f[v]=u);
            sz[u]+=sz[v];
        }
    }
}
void dfs2(int u,int fa){
    for(int v:g[u]){
        if(v^fa){
            if((sz[v]<<1)>sz[u]){
                t[h[u]=v]=t[u];
            }else{
                t[v]=v;
            }
            dfs2(v,u);
        }
    }
}
void dfs3(int u,int fa){
    dfn[u]=++dfn[0];
    if(h[u]){
        dfs3(h[u],u);
    }
    for(int v:g[u]){
        if((v^fa)&&(v^h[u])){
            dfs3(v,u);
        }
    }
}
void add(node*&x,int l,int r,int k){
    if(x==NULL){
        x=new node;
    }
    if(l^r){
        int m=(l+r)>>1;
        if(k<=m){
            add(ls,l,m,k);
        }else{
            add(rs,m+1,r,k);
        }
        x->v=0;
        if(ls!=NULL){
            x->v|=ls->v;
        }
        if(x->v){
            return;
        }
        if(rs!=NULL){
            x->v|=rs->v;
        }
    }else{
        x->v=1;
    }
}
bool query(node*x,int l,int r,int ql,int qr){
    if(ql<=l&&r<=qr){
        return x->v;
    }
    int m=(l+r)>>1;
    bool v=0;
    if(ql<=m&&ls!=NULL){
        v|=query(ls,l,m,ql,qr);
    }
    if(v){
        return 1;
    }
    if(qr>m&&rs!=NULL){
        v|=query(rs,m+1,r,ql,qr);
    }
    return v;
}
bool chainquery(int x,int y,int c){
    bool v=0;
    while(t[x]^t[y]){
        if(d[t[x]]<d[t[y]]){
            swap(x,y);
        }
        v|=query(rt[c],1,n,dfn[t[x]],dfn[x]);
        if(v){
            return 1;
        }
        x=f[t[x]];
    }
    if(d[x]>d[y]){
        swap(x,y);
    }
    return query(rt[c],1,n,dfn[x],dfn[y]);
}
int main(){
    while(~scanf("%d%d",&n,&q)){
        for(int i=1;i<=n;++i){
            scanf("%d",a+i);
            g[i].clear();
            rt[i]=NULL;
            h[i]=0;
        }
        rt[dfn[0]=0]=NULL;
        for(int i=1,u,v;i<n;++i){
            scanf("%d%d",&u,&v);
            g[u].push_back(v);
            g[v].push_back(u);
        }
        dfs1(1,0);
        dfs2(t[1]=1,0);
        dfs3(1,0);
        for(int i=1;i<=n;++i){
            add(rt[a[i]],1,n,dfn[i]);
        }
        for(int i=1,u,v,w;i<=q;++i){
            scanf("%d%d%d",&u,&v,&w);
            if(rt[w]==NULL){
                puts("NotFind");
                continue;
            }
            if(chainquery(u,v,w)){
                puts("Find");
            }else{
                puts("NotFind");
            }
        }
    }
}