【bzoj3697】【采药人的路径】【点分治】

Description

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

Input

第1行包含一个整数N。
接下来N-1行，每行包含三个整数a_i、b_i和t_i，表示这条路上药材的类型。

Output

输出符合采药人要求的路径数目。

Sample Input

7
1 2 0
3 1 1
2 4 0
5 2 0
6 3 1
5 7 1

Sample Output

1

HINT

对于100%的数据，N ≤ 100,000。

题解：首先点分治是肯定的。

            我们把种类看成边权，并把0看成-1.

            对于一棵分治出来的树的子树。

            用f[i][0/1],g[i][0/1] 分别表示当前子树,当前子树之前的子树中权值为i的从子树根节点开始的路径条数.

            0/1表示路径上是否存在一个前缀和为i的点。

            那么每部分的答案就是g[0][0]*f[0][0]+sigma(g[-i][0]*f[i][1]+g[-i][1]*f[i][0]+g[-i][1]*f[i][1])

            为了好处理,我们可以把路径的权值加上n.

代码:

#include<iostream>
#include<cstdio>
#include<cstring>
#define N 100010
using namespace std;
int mx[N],point[N],cnt,next[N*2],deep[N],dis[N],rt,s,size[N],n,x,y,v,p[N*2],mt;
long long ans,g[N*2][2],f[N*2][2];
bool vis[N];
struct use{int st,en,v;}e[N*2];
void add(int x,int y,int v){
  next[++cnt]=point[x];point[x]=cnt;e[cnt].en=y;e[cnt].v=v;
}
void getroot(int x,int fa){
   size[x]=1;mx[x]=0;
   for (int i=point[x];i;i=next[i])
    if (!vis[e[i].en]&&e[i].en!=fa){
      getroot(e[i].en,x);size[x]+=size[e[i].en];
      mx[x]=max(mx[x],size[e[i].en]);
    }	
   mx[x]=max(mx[x],s-size[x]);
   if (mx[x]<mx[rt]) rt=x;
}
void dfs(int x,int fa){
  mt=max(mt,deep[x]);
  if (p[dis[x]]) f[dis[x]][1]++;
  else f[dis[x]][0]++;
  p[dis[x]]++;
  for (int i=point[x];i;i=next[i])
   if (!vis[e[i].en]&&e[i].en!=fa){
   	 deep[e[i].en]=deep[x]+1;
   	 dis[e[i].en]=dis[x]+e[i].v;
   	 dfs(e[i].en,x);
   }	
  p[dis[x]]--;
}
void solve(int x){
  int mx=0;
  vis[x]=1;g[n][0]=1;
  for (int i=point[x];i;i=next[i])
   if (!vis[e[i].en]){
   	 dis[e[i].en]=e[i].v+n;deep[e[i].en]=1;
   	 mt=1;dfs(e[i].en,0);mx=max(mx,mt);
   	 ans+=(g[n][0]-1)*f[n][0];
	 for (int j=-mt;j<=mt;j++)
   	   ans+=g[n-j][1]*f[n+j][0]+g[n-j][0]*f[n+j][1]+g[n-j][1]*f[n+j][1]; 
     for (int j=-mt+n;j<=mt+n;j++){
     	g[j][0]+=f[j][0];
     	g[j][1]+=f[j][1];
     	f[j][0]=f[j][1]=0;
     }
   }
   for (int i=-mx+n;i<=mx+n;i++) g[i][0]=g[i][1]=0;
   for (int i=point[x];i;i=next[i])
    if (!vis[e[i].en]){
       s=size[e[i].en];
       rt=0;getroot(e[i].en,0);
       solve(rt); 
    }
}
int main(){
  scanf("%d",&n);
  for (int i=1;i<=n-1;i++){
    scanf("%d%d%d",&x,&y,&v);
    if (!v) v=-1;add(x,y,v);add(y,x,v); 
  }	
  mx[0]=s=n;
  getroot(1,0);
  solve(rt);cout<<ans<<endl;
}