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

题解
点分治
这样我们枚举根节点的每个子树。用f[i][0…1]，g[i][0…1]分别表示前面几个子树以及当前子树和为i的路径数目，0和1用于区分路径上是否存在前缀和为i的节点。那么当前子树的贡献就是f[0][0] * g[0][0] + Σf [i][0] * g [-i][1] + f[i][1] * g[-i][0] + f[i][1] * g[-i][1]，其中i的范围[-d,d]，d为当前子树的深度。

代码

#include<bits/stdc++.h>
#define N 500005
#define ll long long
using namespace std;
inline int read()
{
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
int n,sum,rt,mxdeep,tot;
bool flag[200005];
int t[200005],mx[100005],size[100005],dep[100005],dis[100005];
int Head[100005],ret[200005],len[200005],Next[200005];
ll ans,g[200005][2],f[200005][2];
inline void ins(int u,int v,int l)
{
    ret[++tot]=v;len[tot]=l;
    Next[tot]=Head[u];Head[u]=tot;
}
void getroot(int u,int f)
{
    size[u]=1;mx[u]=0;
    for (int i=Head[u];i;i=Next[i])
    {
        int v=ret[i];
        if (v!=f&&!flag[v])
        {
            getroot(v,u);
            size[u]+=size[v];
            mx[u]=max(mx[u],size[v]);
        }
    }
    mx[u]=max(mx[u],sum-size[u]);
    if (mx[u]<mx[rt]) rt=u;
}
void dfs(int u,int pre)
{
    mxdeep=max(mxdeep,dep[u]);
    if (t[dis[u]])f[dis[u]][1]++;
    else f[dis[u]][0]++;
    t[dis[u]]++;
    for (int i=Head[u];i;i=Next[i])
    {
        int v=ret[i];
        if (v!=pre&&!flag[v])
        {
            dep[v]=dep[u]+1;
            dis[v]=dis[u]+len[i];
            dfs(v,u);
        }
    }
    t[dis[u]]--;
}
void solve(int x)
{
    int mx=0;
    flag[x]=1;g[n][0]=1;
    for (int i=Head[x];i;i=Next[i])
    {
        int v=ret[i];
        if (!flag[v])
        {
            dis[v]=n+len[i];dep[v]=1;
            mxdeep=1;dfs(v,x);mx=max(mxdeep,mx);
            ans+=(g[n][0]-1)*f[n][0];
            for (int j=-mxdeep;j<=mxdeep;j++)
                ans+=g[n-j][0]*f[n+j][1]+g[n-j][1]*f[n+j][0]+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 i=Head[x];i;i=Next[i])
    {
        if (!flag[ret[i]])
        {
            sum=size[ret[i]];rt=0;
            getroot(ret[i],0);
            solve(rt);
        }
    }
}
int main()
{
    n=read();
    for (int i=1;i<n;i++)
    {
        int a=read(),b=read(),t=read();
        if (t==0) t=-1;
        ins(a,b,t);ins(b,a,t);
    }
    sum=mx[0]=n;
    getroot(1,0);
    solve(rt);
    printf("%lld",ans);
    return 0;
}