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。

贴个题解：http://hzwer.com/4526.html

先点分治，然后对于重心的孩子节点，用f[i][0…1]和g[i][0…1]分别表示前几棵子树和为i的方案数，用0和1分别表示是否有前缀和为i的方案..

然后就用下面的公式咯：$ans=f_{0,0}*g_{0,0}+\sum f_{i,0}*g_{-i,1}+f_{i,1}*g_{-i,0}+f_{i,1}*g_{-i,1}\ (i\in[-maxdep,maxdep])$

#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <algorithm>
#define LL long long
using namespace std;
const LL Maxn = 100010;
struct node {
    LL x, y, next, d;
}a[Maxn*2]; LL first[Maxn], len;
void ins ( LL x, LL y, LL d ){
    len ++;
    a[len].x = x; a[len].y = y; a[len].d = d;
    a[len].next = first[x]; first[x] = len;
}
LL n, core, ms[Maxn], s[Maxn], alls, ans;
bool done[Maxn];
LL f[Maxn*2][2], g[Maxn*2][2], d[Maxn], t[Maxn*2], dep[Maxn], md;
LL _max ( LL x, LL y ){ return x > y ? x : y; }
void dfs ( LL x, LL fa ){
    s[x] = 1; ms[x] = 0;
    for ( LL k = first[x]; k; k = a[k].next ){
        LL y = a[k].y;
        if ( y == fa || done[y] == true ) continue;
        dfs ( y, x );
        s[x] += s[y];
        if ( s[y] > ms[x] ) ms[x] = s[y];
    }
    if ( alls-s[x] > ms[x] ) ms[x] = alls-s[x];
    if ( ms[x] < ms[core] ) core = x;
}
void clac ( LL x, LL fa ){
    md = _max ( md, dep[x] );
    if ( t[n+d[x]] == 0 ){
        g[n+d[x]][0] ++;
    }
    else g[n+d[x]][1] ++;
    t[n+d[x]] ++;
    for ( LL k = first[x]; k; k = a[k].next ){
        LL y = a[k].y;
        if ( y == fa || done[y] == true ) continue;
        d[y] = d[x]+a[k].d; dep[y] = dep[x]+1;
        clac ( y, x );
    }
    t[n+d[x]] --;
}
void solve ( LL x ){
    done[x] = true;
    LL mx = 0;
    f[n][0] = 1;
    for ( LL k = first[x]; k; k = a[k].next ){
        LL y = a[k].y;
        if ( done[y] == true ) continue;
        md = 0;
        d[y] = a[k].d; dep[y] = 1;
        clac ( y, x );
        mx = _max ( mx, md );
        ans += (f[n][0]-1)*g[n][0];
        for ( LL p = -md; p <= md; p ++ ){
            ans += f[n+p][1]*g[n-p][0]+f[n+p][0]*g[n-p][1]+f[n+p][1]*g[n-p][1];
        }
        for ( LL p = -md; p <= md; p ++ ){
            f[n+p][0] += g[n+p][0];
            f[n+p][1] += g[n+p][1];
            g[n+p][0] = g[n+p][1] = 0;
        }
    }
    for ( LL p = -mx; p <= mx; p ++ ){
        f[n+p][0] = f[n+p][1] = 0;
    }
    for ( LL k = first[x]; k; k = a[k].next ){
        LL y = a[k].y;
        if ( done[y] == true ) continue;
        core = 0;
        ms[0] = alls = s[y];
        dfs ( y, 0 );
        solve (core);
    }
}
int main (){
    LL i, j, k;
    scanf ( "%lld", &n );
    for ( i = 1; i < n; i ++ ){
        LL x, y, dd;
        scanf ( "%lld%lld%lld", &x, &y, &dd );
        if ( dd == 0 ) dd --;
        ins ( x, y, dd ); ins ( y, x, dd );
    }
    core = 0;
    ms[0] = alls = n;
    dfs ( 1, 0 );
    memset ( done, false, sizeof (done) );
    ans = 0;
    solve (core);
    printf ( "%lld\n", ans );
    return 0;
}