[BZOJ3697] 采药人的路径

传送门

http://www.lydsy.com/JudgeOnline/problem.php?id=3697

题目大意

给定一棵树，找一些路径满足，路径上0,1数量相等，并在路径上找到一个点(休息站)，改点到路径两端上0,1数量也相等，同一条路径点不同，记为不同，询问有多少条路径满足条件

题解

点分治
对于一个点$u$我们遍历每棵子树中的节点$v$，求出$dis(u,v)$，我们记录这条路径上有无节点$t$使得$dis(v,t)=0$存在数组里
休息站一定在起点到根的路径上或根到终点的路径上
合并两条路径时选取只有一条有休息站的或者两条都有的

const
    maxn=200005;
var
    w:array[0..3*maxn,1..3]of longint;
    size,maxx,vis,dis:array[0..maxn]of longint;
    x:array[-maxn..maxn]of longint;
    f,g:array[-maxn..maxn,0..1]of int64;
    i,j,k:longint;
    n,m,len,a,b,c,root,cnt,mx,mxdep:longint;
    ans:int64;
function max(a,b:longint):longint;
begin if a>b then exit(a) else exit(b); end;

procedure init(a,b,c:longint);
begin
    w[len,1]:=b; w[len,2]:=c; w[len,3]:=0;
    if (w[a,3]=0) then w[a,3]:=len else w[w[a,1],3]:=len;
    w[a,1]:=len; inc(len); 
end;

procedure getroot(a,fa:longint);
var tt:longint;
begin
    tt:=w[a,3]; size[a]:=1; maxx[a]:=0;
    while (tt<>0) do
    begin
        if (w[tt,1]<>fa)and(vis[w[tt,1]]=0) then 
        begin
            getroot(w[tt,1],a);
            inc(size[a],size[w[tt,1]]);
            maxx[a]:=max(maxx[a],size[w[tt,1]]);
        end;
        tt:=w[tt,3];
    end;
    maxx[a]:=max(maxx[a],cnt-size[a]);
    if (root=0)or(maxx[a]<maxx[root]) then root:=a;
end;

procedure getdep(a,fa,b:longint);
var tt:longint;
begin
    tt:=w[a,3]; mxdep:=max(mxdep,b);
    if (x[dis[a]]=0) then g[dis[a],0]:=g[dis[a],0]+1 else g[dis[a],1]:=g[dis[a],1]+1; 
    inc(x[dis[a]]);
    while (tt<>0) do
    begin
        if (w[tt,1]<>fa)and(vis[w[tt,1]]=0) then
        begin
            dis[w[tt,1]]:=dis[a]+w[tt,2];
            getdep(w[tt,1],a,b+1);
        end;
        tt:=w[tt,3];
    end;
    dec(x[dis[a]]);
end;

procedure treedc(a:longint);
var tt,i:longint;
begin
    tt:=w[a,3]; vis[a]:=1; mx:=0; f[0,0]:=1;
    while (tt<>0) do
    begin
        if (vis[w[tt,1]]=0) then 
        begin
            mxdep:=0;
            dis[w[tt,1]]:=w[tt,2]; 
            mxdep:=1; getdep(w[tt,1],0,1); mx:=max(mx,mxdep);
            ans:=ans+g[0,0]*(f[0,0]-1);
            for i:=-mxdep to mxdep do
                ans:=ans+g[i,0]*f[-i,1]+g[i,1]*f[-i,0]+g[i,1]*f[-i,1];
            for i:=-mxdep to mxdep do
            begin
                f[i,0]:=f[i,0]+g[i,0]; f[i,1]:=f[i,1]+g[i,1];
                g[i,0]:=0; g[i,1]:=0;
            end;
        end;
        tt:=w[tt,3];
    end;
    for i:=-mx to mx do
    begin f[i,0]:=0; f[i,1]:=0; end;
    tt:=w[a,3];
    while (tt<>0) do
    begin
        if (vis[w[tt,1]]=0) then 
        begin
            root:=0; cnt:=size[w[tt,1]]; getroot(w[tt,1],0);
            treedc(root);
        end;
        tt:=w[tt,3];
    end;
end;

begin
    readln(n); len:=n+1;
    for i:=1 to n-1 do
    begin readln(a,b,c); if (c=0) then c:=-1; init(a,b,c); init(b,a,c); end;
    root:=0; cnt:=n; getroot(1,0); 
    ans:=0; treedc(root);
    writeln(ans);
end.