dsu on tre模板
阿新 • • 發佈:2020-11-22
/* 詢問每個節點子樹上的顏色數 */ #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]); } }