计蒜客 - 子树权值计数(DFS序+莫队)

题目传送门
先DFS一遍,把每个以 u (1<=u<=n) 为根的子树的最小DFS序和最大DFS序处理出来,再用莫队即可
具体看代码:

#include<bits/stdc++.h>
using namespace std;

const int maxn=100000+100;

struct Edge{

    int to,next;
}edge[maxn<<1];
int head[maxn],tot;
int val[maxn];

int Id[maxn],Rev[maxn],id,Max[maxn];
int n,k,m;
struct Mo{

    int l,r;
    int id;
}mo[maxn];
int R[maxn],Ans[maxn],block,key;
inline bool cmp(Mo a,Mo b){

    return R[a.l]==R[b.l]?a.r<b.r:R[a.l]<R[b.l];
}

void DFS(int u,int fa){

    Id[u]= ++id;
    Rev[id]=u;
    Max[u]=id;
    for(int i=head[u];i!=-1;i=edge[i].next){

        Edge e=edge[i];
        int v=e.to;
        if(v==fa) continue;
        DFS(v,u);
        Max[u]=max(Max[u],Max[v]);
    }
}

unordered_map<int,int>M;
inline void kaven(int ind,int v){

    int cnt=(M[val[Rev[ind]]]+=v);
    if(cnt==k) key++;
    if(v==1&&cnt==k+1) key--;
    if(v==-1&&cnt==k-1) key--;
}

int main(){

    scanf("%d%d",&n,&k);
    block=(int)sqrt(1.0*n);
    for(int i=1;i<=n;i++) scanf("%d",&val[i]),R[i]=i/block,head[i]=-1;
    for(int i=1;i<n;i++){

        int u,v;
        scanf("%d%d",&u,&v);
        edge[tot].to=v;
        edge[tot].next=head[u];
        head[u]=tot++;
        edge[tot].to=u;
        edge[tot].next=head[v];
        head[v]=tot++;
    }
    id=0;
    DFS(1,-1);
    scanf("%d",&m);
    for(int i=1,root;i<=m;i++) scanf("%d",&root),mo[i].l=Id[root],mo[i].r=Max[root],mo[i].id=i;
    sort(mo+1,mo+1+m,cmp);
    int l=1,r=0;
    for(int i=1;i<=m;i++){

        while(l>mo[i].l) kaven(l-1,1),l--;
        while(r<mo[i].r) kaven(r+1,1),r++;
        while(l<mo[i].l) kaven(l,-1),l++;
        while(r>mo[i].r) kaven(r,-1),r--;

        Ans[mo[i].id]=key;
    }
    for(int i=1;i<=m;i++) printf("%d\n",Ans[i]);
}

猜你喜欢

转载自blog.csdn.net/qq_37960603/article/details/81700032