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