【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。

Source

WA了，挖个坑，以后再说

代码：

#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
using namespace std;

const int SZ = 2000010;
const int INF = 1000000010;

int head[SZ],nxt[SZ],tot = 0,n;

struct edge{
int t,d;
}l[SZ];

void build(int f,int t,int d)
{
l[++ tot].t = t;
l[tot].d = d;
nxt[tot] = head[f];
head[f] = tot;
}

int ans = 0,maxn = INF,root;

bool rt[SZ];

int find(int u,int fa,int n)
{
int sz = 1;
int now = 0;
for(int i = head[u];i;i = nxt[i])
{
int v = l[i].t;
if(v != fa && !rt[v])
{
int son = find(v,u,n);
sz += son;
now = max(now,son);
}
}
now = max(now,n - sz);
if(maxn > now) maxn = now,root = u;
return sz;
}

int dist[SZ],deep[SZ];

int dfsdist(int u,int fa,int d)
{
int ans = 0;
if(d == 0) ans = 1;
for(int i = head[u];i;i = nxt[i])
{
int v = l[i].t;
if(!rt[v] && v != fa)
{
ans += dfsdist(v,u,l[i].d + d);
}
}
return ans;
}

int dfssz(int u,int fa)
{
int sz = 1;
for(int i = head[u];i;i = nxt[i])
{
int v = l[i].t;
if(v != fa && !rt[v])
sz += dfssz(v,u);
}
return sz;
}

void dfs(int x,int fa)
{
int sz = dfssz(x,fa);
maxn = INF;
find(x,fa,sz);
int u = root;
//  cout<<u<<endl;
int tot = 0;
rt[u] = 1;
for(int i = head[u];i;i = nxt[i])
{
int v = l[i].t;
if(!rt[v])
{
int tmp = dfsdist(v,u,l[i].d);
ans += tmp * tot;
tot += tmp;
dfs(v,u);
}
}
}

void scanf(int &n)
{
n = 0;
char a = getchar();
bool flag = 0;
while(a < '0' || a > '9') { if(a == '-') flag = 1; a = getchar(); }
while(a >= '0' && a <= '9') n = (n << 3) + (n << 1) + a - '0',a = getchar();
if(flag) n = -n;
}

int main()
{
freopen("in.txt","r",stdin);
freopen("out.txt","w",stdout);
scanf(n);
for(int i = 1;i < n;i ++)
{
int a,b,c;
scanf(a),scanf(b),scanf(c);
if(!c) c = -1;
build(a,b,c);   build(b,a,c);
}
dfs(1,0);
printf("%d",ans);
return 0;
}