【BZOJ3697】采药人的路径（点分治）

题意：采药人的药田是一个树状结构，每条路径上都种植着同种药材。
采药人以自己对药材独到的见解，对每种药材进行了分类。大致分为两类，一种是阴性的，一种是阳性的。
采药人每天都要进行采药活动。他选择的路径是很有讲究的，他认为阴阳平衡是很重要的，所以他走的一定是两种药材数目相等的路径。采药工作是很辛苦的，所以他希望他选出的路径中有一个可以作为休息站的节点（不包括起点和终点），满足起点到休息站和休息站到终点的路径也是阴阳平衡的。他想知道他一共可以选择多少种不同的路径。

n<=100000

思路：RYZ作业

From hzwer：

来自出题人hta的题解。。

本题可以考虑树的点分治。问题就变成求过根满足条件的路径数。

路径上的休息站一定是在起点到根的路径上，或者根到终点的路径上。

如何判断一条从根出发的路径是否包含休息站？只要在dfs中记录下这条路径的和x，同时用个标志数组判断这条路径是否存在前缀和为x的节点。

这样我们枚举根节点的每个子树。用f[i][0…1]，g[i][0…1]分别表示前面几个子树以及当前子树和为i的路径数目，0和1用于区分路径上是否存在前缀和为i的节点。那么当前子树的贡献就是f[0][0] * g[0][0] + Σf [i][0] * g [-i][1] + f[i][1] * g[-i][0] + f[i][1] * g[-i][1]，其中i的范围[-d,d]，d为当前子树的深度。

 var head,vet,next,len,d,dis,son,flag,dep:array[..]of longint;
     c:array[..]of longint;
     b:array[-..]of longint;
     f,g:array[-..,..]of int64;
     n,m,i,x,y,z,tot,sum,root,mxdep:longint;
     ans:int64;

 procedure add(a,b,c:longint);
 begin
  inc(tot);
  next[tot]:=head[a];
  vet[tot]:=b;
  len[tot]:=c;
  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
  son[u]:=; c[u]:=;
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if (v<>fa)and(flag[v]=) then
   begin
    getroot(v,u);
    son[u]:=son[u]+son[v];
    c[u]:=max(c[u],son[v]);
   end;
   e:=next[e];
  end;
  c[u]:=max(c[u],sum-c[u]);
  if c[u]<c[root] then root:=u;
 end;

 procedure dfs(u,fa:longint);
 var e,v:longint;
 begin
  mxdep:=max(mxdep,dep[u]);
  if b[dis[u]]> then inc(f[dis[u],])
   else inc(f[dis[u],]);
  inc(b[dis[u]]);
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if (v<>fa)and(flag[v]=) then
   begin
    dep[v]:=dep[u]+;
    dis[v]:=dis[u]+len[e];
    dfs(v,u);
   end;
   e:=next[e];
  end;
  dec(b[dis[u]]);
 end;

 procedure solve(u:longint);
 var e,v,i,mx:longint;
 begin
  mx:=;
  flag[u]:=; g[,]:=;
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if flag[v]= then
   begin
    dis[v]:=len[e];
    dep[v]:=; mxdep:=;
    dfs(v,u);
    mx:=max(mx,mxdep);
    ans:=ans+(g[,]-)*f[,];
    for i:=-mxdep to mxdep do
     ans:=ans+g[-i,]*f[i,]+g[-i,]*f[i,]+g[-i,]*f[i,];
    for i:=-mxdep to mxdep do
    begin
     g[i,]:=g[i,]+f[i,];
     g[i,]:=g[i,]+f[i,];
     f[i,]:=; f[i,]:=;
    end;
   end;
   e:=next[e];
  end;
  for i:=-mx to mx do
  begin
   g[i,]:=; g[i,]:=;
  end;
  e:=head[u];
  while e<> do
  begin
   v:=vet[e];
   if flag[v]= then
   begin
    sum:=son[v]; root:=;
    getroot(v,);
    solve(root);
   end;
   e:=next[e];
  end;
 end;

 begin
  assign(input,'bzoj3697.in'); reset(input);
  assign(output,'bzoj3697.out'); rewrite(output);
  readln(n);
  for i:= to n- do
  begin
   readln(x,y,z);
   if z= then z:=-;
   add(x,y,z);
   add(y,x,z);
  end;
  root:=; sum:=n; c[]:=n;
  getroot(,);
  solve(root);
  writeln(ans);
  close(input);
  close(output);
 end.