BZOJ 3127 [Usaco2013 Open]Yin and Yang（树点分治）

【题目链接】 http://www.lydsy.com/JudgeOnline/problem.php?id=3127

【题目大意】

　　给出一棵01边权树，求存在多少条路径，使得路径上0和1的数量相同，
　　并且在路劲中能找到至少一个中断点，使得分为两段01数量相同的路径

【题解】

　　我们对这棵树进行点分治，每次只考虑经过重心的路径，
　　我们将路径权值和分出现一次和出现多次进行统计，如果出现一次，
　　则在之前出现了多次的权值数组中查相反数，如果出现多次，
　　则是一次和多次的权值数组的和，权值和为0的特殊情况则另需考虑以重心为端点的路径。
　　每个子树先计算答案，然后统计入权值数组。

【代码】

#include <cstdio>
#include <algorithm>
#include <vector>
#include <cstring>
using namespace std;
const int N=200020; 
vector<int> v[N],e[N];
typedef long long LL;
const int base=100010;
int n,cnt,root,sum,nowT=0;
int mark[N],T1[N],T2[N],size[N],d[N],dp[N],t[N],vis[N];
LL S1[N],S2[N],ans;
void getroot(int x,int fx){
    size[x]=1; dp[x]=0;
    for(int i=0;i<v[x].size();i++){
        int y=v[x][i];
        if(!vis[y]&&y!=fx){
            getroot(y,x);
            size[x]+=size[y];
            dp[x]=max(dp[x],size[y]);
        }
    }dp[x]=max(dp[x],sum-size[x]);
    if(dp[x]<dp[root])root=x;
}
void getdeep(int x,int fx){
    if(mark[base+d[x]]){  
        mark[base+d[x]]++;
        if(!d[x])ans++;
        if(T2[base+d[x]]==root)S2[base+d[x]]++;
        else T2[base+d[x]]=root,S2[base+d[x]]=1;
    }else{
        mark[base+d[x]]++;
        if(T1[base+d[x]]==root)S1[base+d[x]]++;
        else T1[base+d[x]]=root,S1[base+d[x]]=1;
    } 
    for(int i=0;i<v[x].size();i++){
        int y=v[x][i],w=e[x][i];
        if(!vis[y]&&y!=fx){
            d[y]=d[x]+w;
            getdeep(y,x);
        }
    }mark[base+d[x]]--;
}
void caldeep(int x,int fx){
    if(mark[base-d[x]]){
    	mark[base-d[x]]++;
        if(T2[base-d[x]]==root)ans+=S2[base-d[x]];
        else T2[base-d[x]]=root,S2[base-d[x]]=0;
        if(T1[base-d[x]]==root)ans+=S1[base-d[x]];
        else T1[base-d[x]]=root,S1[base-d[x]]=0;
    }else{
        mark[base-d[x]]++;
        if(T2[base-d[x]]==root)ans+=S2[base-d[x]];
        else T2[base-d[x]]=root,S2[base-d[x]]=0;
        if(T1[base-d[x]]==root){
        	if(!d[x])ans+=S1[base-d[x]];
        }else T1[base-d[x]]=root,S1[base-d[x]]=0;
    }
    for(int i=0;i<v[x].size();i++){
        int y=v[x][i],w=e[x][i];
        if(!vis[y]&&y!=fx){
            d[y]=d[x]+w;
            caldeep(y,x);
        }
    }mark[base-d[x]]--;
}
void cal(int x){
    for(int i=0;i<v[x].size();i++){
        int y=v[x][i],w=e[x][i];
        if(!vis[y]){
            d[y]=w;
            caldeep(y,x);
            getdeep(y,x);
        }
    }
}
void solve(int x){
    cal(x); vis[x]=1;
    for(int i=0;i<v[x].size();i++){
        int y=v[x][i];
        if(!vis[y]){
            root=0;sum=size[y];
            getroot(y,0);
            solve(root);
        }
    }
}
int main(){
    while(~scanf("%d",&n)){
        ans=0;
        for(int i=1;i<=n;i++)v[i].clear(),e[i].clear();
        for(int i=1;i<n;i++){
            int x,y,w;
            scanf("%d%d%d",&x,&y,&w); 
            v[x].push_back(y);
            v[y].push_back(x);
            e[x].push_back(w?1:-1);
            e[y].push_back(w?1:-1);
        }memset(vis,0,sizeof(vis)); 
        dp[0]=sum=n;
        getroot(1,0);
        solve(root);
        printf("%lld
",ans);
    }return 0;
}