【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.
null