【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;
}