U41492 树上数颜色

dsu on tree

#include <bits/stdc++.h>
#define inf 2333333333333333
#define N 1000010
#define p(a) putchar(a)
#define For(i,a,b) for(int i=a;i<=b;++i)
//by war
//2020.8.3
using namespace std;
int n,x,y,m;
int c[N],son[N],sz[N],ans[N],sum;
bool vis[N];
struct node{
    int n;
    node *next;
}*e[N];

void in(int &x){
    int y=1;char c=getchar();x=0;
    while(c<'0'||c>'9'){if(c=='-')y=-1;c=getchar();}
    while(c<='9'&&c>='0'){ x=(x<<1)+(x<<3)+c-'0';c=getchar();}
    x*=y;
}
void o(int x){
    if(x<0){p('-');x=-x;}
    if(x>9)o(x/10);
    p(x%10+'0');
}

void push(int x,int y){
    node *p;
    p=new node();
    p->n=y;
    if(e[x]==0)
        e[x]=p;
    else{
        p->next=e[x]->next;
        e[x]->next=p;
    }
}

void clean(int x,int fa){
    vis[c[x]]=0;
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa) continue;
        clean(i->n,x);
    }
}

void dfs1(int x,int fa){
    sz[x]=1;
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa) continue;
        dfs1(i->n,x);
        sz[x]+=sz[i->n];
        if(sz[i->n]>sz[son[x]]) son[x]=i->n;
    }
}

void dfs3(int x,int fa,int so){
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa || i->n==so) continue;
        dfs3(i->n,x,so);
    }
    //todo:计算答案
    if(!vis[c[x]]){
        vis[c[x]]=1;
        sum++;
    }
}

void dfs2(int x,int fa){
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa || i->n==son[x]) continue;
        dfs2(i->n,x);
        clean(i->n,x);
        sum=0;
    }
    if(son[x]) dfs2(son[x],x);
    dfs3(x,fa,son[x]);
    ans[x]=sum;
    return;
}

signed main(){
    in(n);
    For(i,1,n-1){
        in(x);in(y);
        push(x,y);
        push(y,x);
    }
    For(i,1,n) in(c[i]);
    dfs1(1,1);
    dfs2(1,1);
    in(m);
    while(m--){
        in(x);
        o(ans[x]);p('\n');
    }
    return 0;
}