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

Solve

先将标号为0的边改成-1，用f[i][0/1]表示子树的儿子内路径长度为i，之前是否有这样的边；g[i][0/1]则为整个子树。

这样处理后直接点分治。

#include<iostream>#include<cstdlib>#include<cstdio>using namespace std;const int N=100005;int n,fi[N],ne[N<<1],y[N<<1],k[N<<1];int root,maxn[N],sum,son[N],t[N<<1],mxdeep,deep[N];int dis[N];long long ans,f[N<<1][2],g[N<<1][2];bool vis[N];void getroot(int ro,int fa){    maxn[ro]=0;son[ro]=1;    for (int u=fi[ro];u;u=ne[u])        if (y[u]!=fa && !vis[y[u]]){            getroot(y[u],ro);            son[ro]+=son[y[u]];            maxn[ro]=max(maxn[ro],son[y[u]]);        }    maxn[ro]=max(maxn[ro],sum-son[ro]);    if (maxn[ro]<maxn[root])root=ro;}void calc(int ro,int fa){    mxdeep=max(mxdeep,deep[ro]);    if (t[dis[ro]])f[dis[ro]][1]++;    else f[dis[ro]][0]++;    t[dis[ro]]++;    for (int u=fi[ro];u;u=ne[u]){        if (vis[y[u]] || y[u]==fa)continue;        deep[y[u]]=deep[ro]+1;        dis[y[u]]=dis[ro]+k[u];        calc(y[u],ro);    }    t[dis[ro]]--;}void dfs(int ro){    int mx=0;    vis[ro]=1;g[n][0]=1;    for (int u=fi[ro];u;u=ne[u]){        if (vis[y[u]])continue;        dis[y[u]]=n+k[u];deep[y[u]]=1;        mxdeep=1;calc(y[u],0);mx=max(mx,mxdeep);        ans+=(g[n][0]-1)*f[n][0];        for (int j=-mxdeep;j<=mxdeep;++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=n-mxdeep;j<=n+mxdeep;++j){            g[j][0]+=f[j][0];            g[j][1]+=f[j][1];            f[j][0]=f[j][1]=0;        }    }    for (int i=n-mx;i<=n+mx;++i)        g[i][0]=g[i][1]=0;    for (int u=fi[ro];u;u=ne[u])        if (!vis[y[u]]){            sum=son[y[u]];            root=0;            getroot(y[u],0);            dfs(root);        }}int main (){    scanf ("%d",&n);    for (int i=1;i<n;++i){        scanf ("%d%d%d",&y[i+n],&y[i],&k[i]);        if (!k[i])k[i]--;k[i+n]=k[i];        ne[i]=fi[y[i+n]];fi[y[i+n]]=i;        ne[i+n]=fi[y[i]];fi[y[i]]=i+n;    }    sum=maxn[0]=n;    getroot(1,0);    dfs(root);    printf ("%lld",ans);    return 0;}