【ZJOI2017 Round1练习】D8T1 mushroom(点分治)

题意:

思路:

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

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

  1 var head,vet,next:array[1..600000]of longint;
  2     size,f,ans,sun,num,a,flag,vis,b,c:array[0..300000]of longint;
  3     n,m,sum,root,i,tot,now,x,y:longint;
  4 
  5 procedure add(a,b:longint);
  6 begin
  7  inc(tot);
  8  next[tot]:=head[a];
  9  vet[tot]:=b;
 10  head[a]:=tot;
 11 end;
 12 
 13 function max(x,y:longint):longint;
 14 begin
 15  if x>y then exit(x);
 16  exit(y);
 17 end;
 18 
 19 procedure getroot(u,fa:longint);
 20 var e,v:longint;
 21 begin
 22  size[u]:=1; f[u]:=0;
 23  e:=head[u];
 24  while e<>0 do
 25  begin
 26   v:=vet[e];
 27   if (v<>fa)and(flag[v]=0) then
 28   begin
 29    getroot(v,u);
 30    size[u]:=size[u]+size[v];
 31    f[u]:=max(f[u],size[v]);
 32   end;
 33   e:=next[e];
 34  end;
 35  f[u]:=max(f[u],sum-size[u]);
 36  if f[u]<f[root] then root:=u;
 37 end;
 38 
 39 procedure getsize(u,fa:longint);
 40 var e,v:longint;
 41 begin
 42  size[u]:=1;
 43  e:=head[u];
 44  while e<>0 do
 45  begin
 46   v:=vet[e];
 47   if (v<>fa)and(flag[v]=0) then
 48   begin
 49    getsize(v,u);
 50    size[u]:=size[u]+size[v];
 51   end;
 52   e:=next[e];
 53  end;
 54 end;
 55 
 56 procedure dfs(u,fa,k:longint);
 57 var e,v,tmp:longint;
 58 begin
 59  tmp:=0;
 60  ans[u]:=ans[fa];
 61  if vis[a[u]]=0 then
 62  begin
 63   if k=0 then ans[u]:=ans[u]+sum-num[a[u]]
 64    else
 65    begin
 66     if num[a[u]]=0 then
 67     begin
 68      inc(m); b[m]:=a[u];
 69     end;
 70     num[a[u]]:=num[a[u]]+size[u];
 71     ans[now]:=ans[now]+size[u];
 72    end;
 73   vis[a[u]]:=1; tmp:=1;
 74  end;
 75  if k=0 then sun[u]:=sun[u]+ans[u];
 76  e:=head[u];
 77  while e<>0 do
 78  begin
 79   v:=vet[e];
 80   if (v<>fa)and(flag[v]=0) then dfs(v,u,k);
 81   e:=next[e];
 82  end;
 83  if tmp=1 then vis[a[u]]:=0;
 84 end;
 85 
 86 procedure solve(u:longint);
 87 var e,v,i,n:longint;
 88 begin
 89  flag[u]:=1;
 90  n:=0;
 91  e:=head[u];
 92  while e<>0 do
 93  begin
 94   v:=vet[e];
 95   if flag[v]=0 then
 96   begin
 97    inc(n); c[n]:=v;
 98   end;
 99   e:=next[e];
100  end;
101  now:=u;
102  m:=0;
103  vis[a[u]]:=1; ans[u]:=1;
104  sum:=1;
105  for i:=1 to n do
106  begin
107   getsize(c[i],u);
108   dfs(c[i],u,0);
109   dfs(c[i],u,1);
110   sum:=sum+size[c[i]];
111   ans[u]:=ans[u]+size[c[i]];
112  end;
113  sun[u]:=sun[u]+ans[u];
114  for i:=1 to m do num[b[i]]:=0;
115  m:=0;
116  vis[a[u]]:=1; ans[u]:=0;
117  sum:=0;
118  for i:=n downto 1 do
119  begin
120   dfs(c[i],u,0);
121   dfs(c[i],u,1);
122   sum:=sum+size[c[i]];
123   ans[u]:=ans[u]+size[c[i]];
124  end;
125  for i:=1 to m do num[b[i]]:=0;
126  vis[a[u]]:=0;
127  e:=head[u];
128  while e<>0 do
129  begin
130   v:=vet[e];
131   if flag[v]=0 then
132   begin
133    root:=0; sum:=size[v];
134    getroot(v,0);
135    solve(root);
136   end;
137   e:=next[e];
138  end;
139 end;
140 
141 begin
142  assign(input,'mushroom10.in'); reset(input);
143  assign(output,'mushroom10.out'); rewrite(output);
144  readln(n);
145  for i:=1 to n do read(a[i]);
146  for i:=1 to n-1 do
147  begin
148   readln(x,y);
149   add(x,y);
150   add(y,x);
151  end;
152  f[0]:=n; root:=0; sum:=n;
153  getroot(1,0);
154  solve(root);
155  for i:=1 to n do writeln(sun[i]);
156  close(input);
157  close(output);
158 end.
原文地址:https://www.cnblogs.com/myx12345/p/6534565.html