【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[1..210000]of longint;
    c:array[0..110000]of longint;
    b:array[-100000..100000]of longint;
    f,g:array[-100000..100000,0..1]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]:=1; c[u]:=0;
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if (v<>fa)and(flag[v]=0) 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]]>0 then inc(f[dis[u],1])
  else inc(f[dis[u],0]);
 inc(b[dis[u]]);
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if (v<>fa)and(flag[v]=0) then
  begin
   dep[v]:=dep[u]+1;
   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:=0;
 flag[u]:=1; g[0,0]:=1;
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if flag[v]=0 then
  begin
   dis[v]:=len[e];
   dep[v]:=1; mxdep:=1;
   dfs(v,u);
   mx:=max(mx,mxdep);
   ans:=ans+(g[0,0]-1)*f[0,0];
   for i:=-mxdep to mxdep do
    ans:=ans+g[-i,1]*f[i,1]+g[-i,0]*f[i,1]+g[-i,1]*f[i,0];
   for i:=-mxdep to mxdep do
   begin
    g[i,0]:=g[i,0]+f[i,0];
    g[i,1]:=g[i,1]+f[i,1];
    f[i,0]:=0; f[i,1]:=0;
   end;
  end;
  e:=next[e];
 end;
 for i:=-mx to mx do
 begin
  g[i,0]:=0; g[i,1]:=0;
 end;
 e:=head[u];
 while e<>0 do
 begin
  v:=vet[e];
  if flag[v]=0 then
  begin
   sum:=son[v]; root:=0;
   getroot(v,0);
   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:=1 to n-1 do
 begin
  readln(x,y,z);
  if z=0 then z:=-1;
  add(x,y,z);
  add(y,x,z);
 end;
 root:=0; sum:=n; c[0]:=n;
 getroot(1,0);
 solve(root);
 writeln(ans);
 close(input);
 close(output);
end.