【洛谷P2664】 树上游戏 点分治
code:
#include <bits/stdc++.h> #define N 200009 #define ll long long #define setIO(s) freopen(s".in","r",stdin) using namespace std; ll Sum[N]; int n,edges,root,sn; int val[N],hd[N],to[N<<1],nex[N<<1],size[N],mx[N],vis[N],A[N]; inline void add(int u,int v) { nex[++edges]=hd[u],hd[u]=edges,to[edges]=v; } void getroot(int u,int ff) { size[u]=1,mx[u]=0; for(int i=hd[u];i;i=nex[i]) { int v=to[i]; if(v==ff||vis[v]) continue; getroot(v,u); size[u]+=size[v]; mx[u]=max(mx[u],size[v]); } mx[u]=max(mx[u],sn-size[u]); if(mx[u]<mx[root]) root=u; } int ou; ll tmp,tot,bu[N]; map<int,ll>cn[N]; map<int,ll>::iterator it; int dep[N],cnt[N],siz[N]; void getnode(int top,int u,int ff,int cur) { if(!cnt[val[u]]) ++cur; ++cnt[val[u]]; Sum[top]+=(ll)cur; siz[u]=1; for(int i=hd[u];i;i=nex[i]) { int v=to[i]; if(v==ff||vis[v]) continue; getnode(top,v,u,cur); siz[u]+=siz[v]; } --cnt[val[u]]; } void get_col(int top,int u,int ff) { if(!cnt[val[u]]) cn[top][val[u]]+=(ll)siz[u]; ++cnt[val[u]]; for(int i=hd[u];i;i=nex[i]) { int v=to[i]; if(v==ff||vis[v]) continue; get_col(top,v,u); } --cnt[val[u]]; } void calc_v(int u,int ff) { ll tt=bu[val[u]]; tmp=tmp-bu[val[u]]+ou; bu[val[u]]=ou; Sum[u]+=tmp; for(int i=hd[u];i;i=nex[i]) { int v=to[i]; if(vis[v]||v==ff) continue; calc_v(v,u); } tmp=tmp-bu[val[u]]+tt; bu[val[u]]=tt; } void clr(int u,int ff) { cn[u].clear(); bu[val[u]]=0; for(int i=hd[u];i;i=nex[i]) { int v=to[i]; if(v==ff||vis[v]) continue; clr(v,u); } } void calc(int u) { tot=0; getnode(u,u,0,0); for(int i=hd[u];i;i=nex[i]) { int v=to[i]; if(vis[v]) continue; // memset(cnt,0,sizeof(cnt)); get_col(v,v,u); for(it=cn[v].begin();it!=cn[v].end();it++) { tot+=it->second; bu[it->first]+=it->second; } } for(int i=hd[u];i;i=nex[i]) { int v=to[i]; if(vis[v]) continue; tmp=tot; ou=siz[u]-siz[v]; for(it=cn[v].begin();it!=cn[v].end();it++) { bu[it->first]-=it->second; tmp-=it->second; } ll tt=bu[val[u]]; tmp=tmp-bu[val[u]]+ou; bu[val[u]]=ou; calc_v(v,u); bu[val[u]]=tt; for(it=cn[v].begin();it!=cn[v].end();it++) bu[it->first]+=it->second; } clr(u,0); } void dfs(int u) { calc(u); vis[u]=1; for(int i=hd[u];i;i=nex[i]) { int v=to[i]; if(vis[v]) continue; root=0,sn=size[v],getroot(v,u),dfs(root); } } int main() { // setIO("input"); int i,j; scanf("%d",&n); for(i=1;i<=n;++i) scanf("%d",&val[i]), A[i]=val[i]; sort(A+1,A+1+n); for(i=1;i<=n;++i) val[i]=lower_bound(A+1,A+1+n,val[i])-A; for(i=1;i<n;++i) { int u,v; scanf("%d%d",&u,&v),add(u,v),add(v,u); } sn=mx[0]=n,root=0,getroot(1,0),dfs(root); for(i=1;i<=n;++i) printf("%lld\n",Sum[i]); return 0; }