dsu on tre模板

/*
询问每个节点子树上的颜色数
*/
#include <bits/stdc++.h>
using namespace std;
const int maxn=1e5+5;
struct edge{
    int v,next;
}E[maxn<<1];
int head[maxn],tot=0;
void addedge(int u,int v){
    E[++tot].v=v;
    E[tot].next=head[u];
    head[u]=tot;
}
int son[maxn],sz[maxn];
void dfs1(int u,int fa){//预处理重儿子
    sz[u]=1;
    for(int i=head[u];i;i=E[i].next){
        int v=E[i].v;
        if(v!=fa){
            dfs1(v,u);
            if(sz[v]>sz[son[u]])son[u]=v;
            sz[u]+=sz[v];
        }
    }
}
int a[maxn];
int cnt[maxn],ans[maxn],cur;
int Bigson;
void add(int u,int fa,int val){//暴力统计轻子树
    if(!cnt[a[u]])cur++;
    cnt[a[u]]+=val;
    if(!cnt[a[u]])cur--;
    for(int i=head[u];i;i=E[i].next){
        int v=E[i].v;
        if(v!=fa && v!=Bigson){
            add(v,u,val);
        }
    }
}
void dfs2(int u,int fa,int opt){
    for(int i=head[u];i;i=E[i].next){
        int v=E[i].v;
        if(v!=fa && v!=son[u])dfs2(v,u,0);//轻儿子计算完毕后要消除影响
    }
    if(son[u])dfs2(son[u],u,1),Bigson=son[u];//重儿子影响保留
    add(u,fa,1);//暴力统计轻儿子贡献(和u)
    ans[u]=cur;
    Bigson=0;
    if(!opt)add(u,fa,-1);
}
int main () {
    int n;
    scanf("%d",&n);
    for(int i=1;i<=n-1;i++){
        int u,v;
        scanf("%d%d",&u,&v);
        addedge(u,v);
        addedge(v,u);
    }
    for(int i=1;i<=n;i++){
        scanf("%d",&a[i]);
    }
    dfs1(1,0);
    cur=0;
    dfs2(1,0,1);
    int m;scanf("%d",&m);
    for(int i=1;i<=m;i++){
        int x;
        scanf("%d",&x);
        printf("%d\n",ans[x]);
    }
}
posted @ 2020-11-22 16:58  UCPRER  阅读(75)  评论(0编辑  收藏  举报