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

题意：

思路：

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

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

var head,vet,next:array[1..600000]of longint;
    size,f,ans,sun,num,a,flag,vis,b,c:array[0..300000]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]:=1; f[u]:=0;
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if (v<>fa)and(flag[v]=0) 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]:=1;
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if (v<>fa)and(flag[v]=0) 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:=0;
 ans[u]:=ans[fa];
 if vis[a[u]]=0 then
 begin
  if k=0 then ans[u]:=ans[u]+sum-num[a[u]]
   else
   begin
    if num[a[u]]=0 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]]:=1; tmp:=1;
 end;
 if k=0 then sun[u]:=sun[u]+ans[u];
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if (v<>fa)and(flag[v]=0) then dfs(v,u,k);
  e:=next[e];
 end;
 if tmp=1 then vis[a[u]]:=0;
end;

procedure solve(u:longint);
var e,v,i,n:longint;
begin
 flag[u]:=1;
 n:=0;
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if flag[v]=0 then
  begin
   inc(n); c[n]:=v;
  end;
  e:=next[e];
 end;
 now:=u;
 m:=0;
 vis[a[u]]:=1; ans[u]:=1;
 sum:=1;
 for i:=1 to n do
 begin
  getsize(c[i],u);
  dfs(c[i],u,0);
  dfs(c[i],u,1);
  sum:=sum+size[c[i]];
  ans[u]:=ans[u]+size[c[i]];
 end;
 sun[u]:=sun[u]+ans[u];
 for i:=1 to m do num[b[i]]:=0;
 m:=0;
 vis[a[u]]:=1; ans[u]:=0;
 sum:=0;
 for i:=n downto 1 do
 begin
  dfs(c[i],u,0);
  dfs(c[i],u,1);
  sum:=sum+size[c[i]];
  ans[u]:=ans[u]+size[c[i]];
 end;
 for i:=1 to m do num[b[i]]:=0;
 vis[a[u]]:=0;
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if flag[v]=0 then
  begin
   root:=0; sum:=size[v];
   getroot(v,0);
   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:=1 to n do read(a[i]);
 for i:=1 to n-1 do
 begin
  readln(x,y);
  add(x,y);
  add(y,x);
 end;
 f[0]:=n; root:=0; sum:=n;
 getroot(1,0);
 solve(root);
 for i:=1 to n do writeln(sun[i]);
 close(input);
 close(output);
end.