【ZJOI2017 Round1练习】D8T1 mushroom(点分治)
阅读原文时间:2023年07月16日阅读:1

题意:

思路:

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.