【ZJOI2017 Round1练习】D8T1 mushroom（点分治）

题意：

思路：

num[a[u]]表示存在a[u]这个颜色且终点在u子树中的链长总和

ans[i]表示以当前的u为根，前面的子树对i的贡献之和

 var head,vet,next:array[..]of longint;
     size,f,ans,sun,num,a,flag,vis,b,c:array[..]of longint;
     n,m,sum,root,i,tot,now,x,y:longint;

 procedure add(a,b:longint);
 begin
  inc(tot);
  next[tot]:=head[a];
  vet[tot]:=b;
  head[a]:=tot;
 end;

 function max(x,y:longint):longint;
 begin
  if x>y then exit(x);
  exit(y);
 end;

 procedure getroot(u,fa:longint);
 var e,v:longint;
 begin
  size[u]:=; f[u]:=;
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if (v<>fa)and(flag[v]=) then
   begin
    getroot(v,u);
    size[u]:=size[u]+size[v];
    f[u]:=max(f[u],size[v]);
   end;
   e:=next[e];
  end;
  f[u]:=max(f[u],sum-size[u]);
  if f[u]<f[root] then root:=u;
 end;

 procedure getsize(u,fa:longint);
 var e,v:longint;
 begin
  size[u]:=;
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if (v<>fa)and(flag[v]=) then
   begin
    getsize(v,u);
    size[u]:=size[u]+size[v];
   end;
   e:=next[e];
  end;
 end;

 procedure dfs(u,fa,k:longint);
 var e,v,tmp:longint;
 begin
  tmp:=;
  ans[u]:=ans[fa];
  if vis[a[u]]= then
  begin
   if k= then ans[u]:=ans[u]+sum-num[a[u]]
    else
    begin
     if num[a[u]]= then
     begin
      inc(m); b[m]:=a[u];
     end;
     num[a[u]]:=num[a[u]]+size[u];
     ans[now]:=ans[now]+size[u];
    end;
   vis[a[u]]:=; tmp:=;
  end;
  if k= then sun[u]:=sun[u]+ans[u];
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if (v<>fa)and(flag[v]=) then dfs(v,u,k);
   e:=next[e];
  end;
  if tmp= then vis[a[u]]:=;
 end;

 procedure solve(u:longint);
 var e,v,i,n:longint;
 begin
  flag[u]:=;
  n:=;
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if flag[v]= then
   begin
    inc(n); c[n]:=v;
   end;
   e:=next[e];
  end;
  now:=u;
  m:=;
  vis[a[u]]:=; ans[u]:=;
  sum:=;
  for i:= to n do
  begin
   getsize(c[i],u);
   dfs(c[i],u,);
   dfs(c[i],u,);
   sum:=sum+size[c[i]];
   ans[u]:=ans[u]+size[c[i]];
  end;
  sun[u]:=sun[u]+ans[u];
  for i:= to m do num[b[i]]:=;
  m:=;
  vis[a[u]]:=; ans[u]:=;
  sum:=;
  for i:=n downto  do
  begin
   dfs(c[i],u,);
   dfs(c[i],u,);
   sum:=sum+size[c[i]];
   ans[u]:=ans[u]+size[c[i]];
  end;
  for i:= to m do num[b[i]]:=;
  vis[a[u]]:=;
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if flag[v]= then
   begin
    root:=; sum:=size[v];
    getroot(v,);
    solve(root);
   end;
   e:=next[e];
  end;
 end;

 begin
  assign(input,'mushroom10.in'); reset(input);
  assign(output,'mushroom10.out'); rewrite(output);
  readln(n);
  for i:= to n do read(a[i]);
  for i:= to n- do
  begin
   readln(x,y);
   add(x,y);
   add(y,x);
  end;
  f[]:=n; root:=; sum:=n;
  getroot(,);
  solve(root);
  for i:= to n do writeln(sun[i]);
  close(input);
  close(output);
 end.