【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
,point
,cnt,next[N*2],deep
,dis
,rt,s,size
,n,x,y,v,p[N*2],mt;
long long ans,g[N*2][2],f[N*2][2];
bool vis
;
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
[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
[0]-1)*f
[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;
}