BZOJ 3697/3127 采药人的路径 (点分治)

题目大意：

从前有一棵无向树，树上边权均为$0$或$1$，有一个采药人，他认为如果一条路径上边权为$0$和$1$的边数量相等，那么这条路径阴阳平衡。他想寻找一条合法的采药路径，保证阴阳平衡。然后他发现采药很累，于是乎他需要保证这条路径上有一个中转站，路径两个端点到中转站的路径都需要阴阳平衡 $n \leqslant 10^{5}$，求合法路径数

把$0$边边权变成$-1$，易发现如果一条路径阴阳平衡，那么边权总和为$0$

由于是树上路径的计数问题，考虑树分治，每次选中心作为根，设某点$x$到根的路径的边权和为$dis_{x}$

把中转站加进去会发生什么呢？

如果某个点$x$到根的路径上能设置中转站，那么根到$x$的路径上一定存在一点$y$，$dis_{y}=dis_{x}$，即$x$到$y$的路径$dis$为$0$，这个操作可以用桶实现

现在要在根统计答案了，显然每个点分为种情况，到根的路径能设置中转站和不能，用桶分别计数即可，设为$f[i][0]$和$f[i][1]$，表示$dis=i$时的方案数

发现还有$dis<0$的情况，需要额外记录一个数组$g$

那么答案就是$\sum_{i=1} f[i][0]*g[i][1]+f[i][1]*g[i][0]+f[i][1]*g[i][1]$

还要去掉不合法的情况，在当前根的每个子节点依然进行上述方法计数即可

$dis=0$的情况需要单独讨论，所有$dis=0$的路径都能两两匹配。此外，不以根节点为中转站，且$dis=0$的路径也一定合法

细节比较多

 #include <cmath>
 #include <vector>
 #include <cstdio>
 #include <cstring>
 #include <algorithm>
 #define N1 100010
 #define M1 (N1<<1)
 #define ll long long
 #define inf 233333333
 using namespace std;
 
 int gint()
 {
     int ret=,fh=;char c=getchar();
     while(c<''||c>''){if(c=='-')fh=-;c=getchar();}
     while(c>=''&&c<=''){ret=ret*+c-'';c=getchar();}
     return ret*fh;
 }
 int n;
 
 struct Edge{
 int to[M1],nxt[M1],val[M1],head[N1],cte;
 void ae(int u,int v,int w)
 {cte++;to[cte]=v,nxt[cte]=head[u],val[cte]=w,head[u]=cte;}
 }E;
 
 int use[N1],mi,G,S,tsz,ma;
 int sf[N1],sg[N1],ms[N1],sz[N1],dis[N1];
 int f[N1][],g[N1][]; ll ans; int que[N1],tl;
 //int dep[N1],fa[N1],de;
 void dfs_sum(int u,int dad)
 {
     que[++tl]=u;
     if(dis[u]>=){
         if(sf[dis[u]]) f[dis[u]][]++;
         else f[dis[u]][]++;
         sf[dis[u]]++;  ma=max(ma,dis[u]);
         if(u==S) sf[]--,f[][]--;
     }else{
         if(sg[-dis[u]]) g[-dis[u]][]++;
         else g[-dis[u]][]++;
         sg[-dis[u]]++; ma=max(ma,-dis[u]);
     }
     for(int j=E.head[u];j;j=E.nxt[j])
     {
         if(E.to[j]==dad||use[E.to[j]]) continue;
         dis[E.to[j]]=dis[u]+E.val[j];
         dfs_sum(E.to[j],u);
     }
     if(dis[u]>=) sf[dis[u]]--;
     else sg[-dis[u]]--;
 }
 void gra(int u,int dad)
 {
     sz[u]=; ms[u]=;
     for(int j=E.head[u];j;j=E.nxt[j])
     {
         int v=E.to[j];
         if(v==dad||use[v]) continue;
         gra(v,u); sz[u]+=sz[v];
         ms[u]=max(ms[u],sz[v]);
     }
     ms[u]=max(ms[u],tsz-sz[u]);
     if(ms[u]<ms[G]) G=u;
 }
 void clr()
 {
     int x;
     while(tl)
     {
         x=que[tl]; tl--;
         if(dis[x]>=) f[dis[x]][]=f[dis[x]][]=;
         else g[-dis[x]][]=g[-dis[x]][]=;
     }
     sg[]=sf[]=;
 }
 void calc(int u,int type)
 {
     ma=; dfs_sum(u,-);
     ans+=( 1ll*f[][]*(f[][]-)/ + 1ll*f[][]*(f[][]-)/ + 1ll*f[][]*f[][] )*type;
     for(int i=;i<=ma;i++)
         ans+=( 1ll*f[i][]*g[i][] + 1ll*f[i][]*g[i][] + 1ll*f[i][]*g[i][])*type;
     if(type==) ans+=f[][];
     clr();
 }
 void main_dfs(int u)
 {
     int j,v; use[u]=; S=u; dis[u]=; calc(u,);
     for(j=E.head[u];j;j=E.nxt[j])
     {
         v=E.to[j]; if(use[v]) continue;
         G=; tsz=sz[v]; gra(v,u);
         calc(v,-);
         main_dfs(G);
     }
 }
 
 int main()
 {
     //freopen("t2.in","r",stdin);
     int i,x,y,z;
     scanf("%d",&n);
     for(i=;i<n;i++)
     {
         x=gint(),y=gint(),z=gint();
         z=((z)?:-);
         E.ae(x,y,z),E.ae(y,x,z);
     }
     ms[]=tsz=n; G=; gra(,-); gra(G,-);
     main_dfs(G);
     printf("%lld\n",ans);
     return ;
 }