【POJ1741】Tree（点分治）

  1 var head,vet,next,len,son,flag,d:array[1..50000]of longint;
  2     dep,f:array[0..50000]of longint;
  3     n,m,i,x,y,z,tot,ans,sum,root,k:longint;
  4 
  5 procedure swap(var x,y:longint);
  6 var t:longint;
  7 begin
  8  t:=x; x:=y; y:=t;
  9 end;
 10 
 11 procedure add(a,b,c:longint);
 12 begin
 13  inc(tot);
 14  next[tot]:=head[a];
 15  vet[tot]:=b;
 16  len[tot]:=c;
 17  head[a]:=tot;
 18 end;
 19 
 20 function max(x,y:longint):longint;
 21 begin
 22  if x>y then exit(x);
 23  exit(y);
 24 end;
 25 
 26 procedure getroot(u,fa:longint);
 27 var e,v:longint;
 28 begin
 29  son[u]:=1; f[u]:=0;
 30  e:=head[u];
 31  while e<>0 do
 32  begin
 33   v:=vet[e];
 34   if (v<>fa)and(flag[v]=0) then
 35   begin
 36    getroot(v,u);
 37    son[u]:=son[u]+son[v];
 38    f[u]:=max(f[u],son[v]);
 39   end;
 40   e:=next[e];
 41  end;
 42  f[u]:=max(f[u],sum-f[u]);
 43  if f[u]<f[root] then root:=u;
 44 end;
 45 
 46 procedure getdeep(u,fa:longint);
 47 var e,v:longint;
 48 begin
 49  inc(dep[0]); dep[dep[0]]:=d[u];
 50  e:=head[u];
 51  while e<>0 do
 52  begin
 53   v:=vet[e];
 54   if (v<>fa)and(flag[v]=0) then
 55   begin
 56    d[v]:=d[u]+len[e];
 57    getdeep(v,u);
 58   end;
 59   e:=next[e];
 60  end;
 61 end;
 62 
 63 procedure qsort(l,r:longint);
 64 var i,j,mid:longint;
 65 begin
 66  i:=l; j:=r; mid:=dep[(l+r)>>1];
 67  repeat
 68   while mid>dep[i] do inc(i);
 69   while mid<dep[j] do dec(j);
 70   if i<=j then
 71   begin
 72    swap(dep[i],dep[j]);
 73    inc(i); dec(j);
 74   end;
 75  until i>j;
 76  if l<j then qsort(l,j);
 77  if i<r then qsort(i,r);
 78 end;
 79 
 80 function clac(u,now:longint):longint;
 81 var l,r:longint;
 82 begin
 83  d[u]:=now; dep[0]:=0;
 84  getdeep(u,0);
 85  qsort(1,dep[0]);
 86  clac:=0;
 87  l:=1; r:=dep[0];
 88  while l<r do
 89  begin
 90   if dep[l]+dep[r]<=k then begin clac:=clac+r-l; inc(l); end
 91    else dec(r);
 92  end;
 93 end;
 94 
 95 procedure work(u:longint);
 96 var e,v:longint;
 97 begin
 98  ans:=ans+clac(u,0);
 99  flag[u]:=1;
100  e:=head[u];
101  while e<>0 do
102  begin
103   v:=vet[e];
104   if flag[v]=0 then
105   begin
106    ans:=ans-clac(v,len[e]);
107    sum:=son[v];
108    root:=0;
109    getroot(v,root);
110    work(root);
111   end;
112   e:=next[e];
113  end;
114 end;
115 
116 begin
117  assign(input,‘poj1741.in‘); reset(input);
118  assign(output,‘poj1741.out‘); rewrite(output);
119  while not eof do
120  begin
121   fillchar(head,sizeof(head),0);
122   fillchar(flag,sizeof(flag),0);
123   tot:=0;
124   readln(n,k);
125   if (n=0)and(m=0) then break;
126   for i:=1 to n-1 do
127   begin
128    readln(x,y,z);
129    add(x,y,z);
130    add(y,x,z);
131   end;
132   sum:=n; f[0]:=maxlongint; ans:=0;
133   getroot(1,0);
134   work(root);
135   writeln(ans);
136  end;
137  close(input);
138  close(output);
139 end.