U41492 樹上數顏色(dsu on tree)
阿新 • • 發佈:2021-08-07
\(O(n^2)\)顯然不過我們應該優化成\(O(nlogn)\)
採用樹上啟發式合併
仿照樹鏈剖分的思想,對於每一個位置,我們先處理所有的輕兒子,然後處理重兒子,統計當前節點的答案,最後把輕兒子刪掉就可以了。
這樣全域性一個桶就夠用了。
#include<iostream> #include<cstdio> #include<algorithm> #include<cstring> using namespace std; int n; int x,y; int col[1000005]; struct e{ int to; int ne; }ed[3000005]; int p; int head[1000005]; void add(int f,int to){ p++; ed[p].to=to; ed[p].ne=head[f]; head[f]=p; return ; } int siz[2000005]; int ns; int ans[2000005]; int tem; int son[1000005]; int exi[1000005]; void dfs(int no,int fa){ siz[no]++; for(int i=head[no];i;i=ed[i].ne){ if(ed[i].to==fa) continue; dfs(ed[i].to,no); siz[no]+=siz[ed[i].to]; if(siz[ed[i].to]>siz[son[no]]){ son[no]=ed[i].to; } } } void work(int no,int fa,int ke){ if(!exi[col[no]]){ tem++; } exi[col[no]]+=ke; for(int i=head[no];i;i=ed[i].ne){ int v=ed[i].to; if(v==fa||v==ns) continue; work(v,no,ke); } } void dsu(int no,int fa,int f){ for(int i=head[no];i;i=ed[i].ne){ if(ed[i].to==fa||ed[i].to==son[no]) continue; dsu(ed[i].to,no,0); } if(son[no]) { dsu(son[no],no,1); ns=son[no]; } work(no,fa,1); ns=0; ans[no]=tem; if(!f){ work(no,fa,-1); tem=0; } } int main(){ scanf("%d",&n); for(int i=1;i<n;++i){ scanf("%d%d",&x,&y); add(x,y); add(y,x); } for(int i=1;i<=n;++i){ scanf("%d",&col[i]); } dfs(1,0); //cout<<2344; dsu(1,0,1); scanf("%d",&n); for(int i=1;i<=n;++i){ scanf("%d",&x); printf("%d\n",ans[x]); } return 0; }