BZOJ 3697: 采药人的路径 树的点分治

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。

题解

正常点分治，将0的边权改为-1，设两个数组f1,f2,分别代表之前出现过休息点的路径数目以及之前没出现过休息点的路径数目，然后到每一个点统计一下路径权值和v，若-v也出现过那么f1[v]++，否则f2[v]++，最后的答案就是两条路径都有休息点的方案数加上只有一个休息点的方案数。

#include<cstdio>
#include<cstdlib>
#include<iostream>
#include<iomanip>
#include<ctime>
#include<cmath>
#include<cstring>
#include<string>
#include<algorithm>
using namespace std;
struct bian
{
    int l,r,v;
    bool ban;
}a[400000];
int fir[400000];
int nex[400000];
int tot=1;
void add_edge(int l,int r,int v)
{
    a[++tot].l=l;
    a[tot].r=r;
    a[tot].v=v;
    a[tot].ban=false;
    nex[tot]=fir[l];
    fir[l]=tot;
}
int siz[400000];
int fa[400000];
int center;
void get_tree_center(int u,int fro,int zong)
{
    fa[u]=fro;
    siz[u]=1;
    bool pd=true;
    for(int o=fir[u];o;o=nex[o])
    {
        if(a[o].r==fro || a[o].ban) continue;
        get_tree_center(a[o].r,u,zong);
        if(siz[a[o].r]>zong/2) pd=false;
        siz[u]+=siz[a[o].r];
    }
    if(zong-siz[u]>zong/2) pd=false;
    if(pd) center=u;
}
int f[400000][2];
int g[400000][2];
int sum[400000];
int maxx=0,minn=400000;
void dfs(int u,int fro,int v)
{
    maxx=max(maxx,v);
    minn=min(minn,v);
    if(sum[v]) f[v][1]++;
    else f[v][0]++;
    sum[v]++;
    for(int o=fir[u];o;o=nex[o])
    {
        if(a[o].r==fro || a[o].ban) continue;
        dfs(a[o].r,u,v+a[o].v);
    }
    sum[v]--;
}
long long ans=0;
void tree_divide(int u)
{
    int zong=siz[u];
    get_tree_center(u,0,siz[u]);
    u=center;
    if(fa[u]) siz[fa[u]]=zong-siz[u];
    g[200000][0]=1;
    int midmaxx=0,midminn=400000;
    for(int o=fir[u];o;o=nex[o])
    {
        if(a[o].ban) continue;
        maxx=0,minn=400000;
        dfs(a[o].r,u,200000+a[o].v);
        midmaxx=max(midmaxx,maxx);
        midminn=min(midminn,minn);
        ans+=1ll*(g[200000][0]-1)*f[200000][0];
        for(int i=minn;i<=maxx;i++)
            ans+=1ll*f[i][0]*g[400000-i][1]+1ll*f[i][1]*g[400000-i][0]+1ll*f[i][1]*g[400000-i][1];
        for(int i=minn;i<=maxx;i++) g[i][0]+=f[i][0],g[i][1]+=f[i][1],f[i][0]=f[i][1]=0;
    }
    for(int i=midminn;i<=midmaxx;i++) g[i][0]=g[i][1]=0;
    g[200000][0]=0;
    for(int o=fir[u];o;o=nex[o])
    {
        if(a[o].ban) continue;
        a[o].ban=true;
        a[o^1].ban=true;
        tree_divide(a[o].r);
    }
}
int main()
{
    int n;
    scanf("%d",&n);
    for(int i=1;i<n;i++)
    {
        int l,r,v;
        scanf("%d%d%d",&l,&r,&v);
        if(!v) v=-1;
        add_edge(l,r,v);
        add_edge(r,l,v);
    }
    siz[1]=n;
    tree_divide(1);
    cout<<ans;
    return 0;
}