Boatherds POJ - 2114 （树分治）

思路：还是和以前求路径值<=k，一样的套路，改下cal函数里计算==k的方法就行。

代码：

#include<iostream>
#include<map>
#include<string>
#include<cstring>
#include<vector>
#include<algorithm>
#include<set>
#include<sstream>
#include<cstdio>
#include<cmath>
#include<climits>
using namespace std;
const int maxn=1e4+7;
const int inf=0x3f3f3f3f;
typedef long long ll;
const int mod=1e9+7;
int n,k,allnode;
int head[maxn*2];
int num;
int dp[maxn];
int size[maxn];
int Focus,M;
ll dist[maxn];
int deep[maxn];
bool vis[maxn];
int ans;
struct Edge
{
    int u,v,w,next;
}edge[maxn<<2];
void addEdge(int u,int v,int w)
{
    edge[num].u=u;
    edge[num].v=v;
    edge[num].w=w;
    edge[num].next=head[u];
    head[u]=num++;
}
void init()
{
    memset(head,-1,sizeof(head));
    memset(dist,0,sizeof(dist));
    num=0;
}
void getFocus(int u,int pre)
{
    size[u]=1;
    dp[u]=0;
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==pre||vis[v]) continue;
        getFocus(v,u);
        size[u]+=size[v];
        dp[u]=max(dp[u],size[v]);
    }
    dp[u]=max(dp[u],allnode-size[u]);
    if(M>dp[u])
    {
        M=dp[u];
        Focus=u;
    }
}
void dfs(int u,int pre)
{
    deep[++deep[0]]=dist[u];
    for(int i=head[u];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(v==pre||vis[v]) continue;
        dist[v]=dist[u]+edge[i].w;
        dfs(v,u);
    }
}
int cal(int x,int now)
{
    dist[x]=now,deep[0]=0;
    dfs(x,0);
    int ans=0;
    sort(deep+1,deep+1+deep[0]);
    int l=1,r=deep[0];
    while(l<r)
    {
        if(deep[l]+deep[r]>k)
        {
            r--;
        }
        else if(deep[l]+deep[r]<k)
        {
            l++;
        }
        else
        {
            if(deep[l]==deep[r])
            {
                ans+=(r-l+1)*(r-l)/2;
                break;
            }
            else
            {
                int i=l,j=r;
                while(deep[i]==deep[l]) i++;
                while(deep[j]==deep[r]) j--;
                ans+=(i-l)*(r-j);
                l=i,r=j;
            }
        }
    }
    return ans;
}
void solve(int x)
{
    vis[x]=1;
    ans+=cal(x,0);
    for(int i=head[x];i!=-1;i=edge[i].next)
    {
        int v=edge[i].v;
        if(vis[v]) continue;
        ans-=cal(v,edge[i].w);
        allnode=size[v],Focus=0,M=1e9;
        getFocus(v,x);
        solve(Focus);
    }

}
int main(int argc, char const *argv[])
{
    #ifndef ONLINE_JUDGE
        freopen("in.txt","r",stdin);
        freopen("out.txt","w",stdout);
    #endif
    while(scanf("%d",&n)!=EOF&&n)
    {
        init();
        int u,v,w;
        for(int i=1;i<=n;i++)
        {
            while(scanf("%d",&v)!=EOF&&v)
            {
                scanf("%d",&w);
                addEdge(i,v,w);
                addEdge(v,i,w);
            }
        }
        while(scanf("%d",&k)!=EOF&&k)
        {
            memset(vis,0,sizeof(vis));
            allnode=n,Focus=0,M=1e9;
            ans=0;
            getFocus(1,0);
            solve(Focus);
            if(ans==0)
            {
                puts("NAY");
            }
            else
            {
                puts("AYE");
            }
        }
        printf(".\n");
    }
    return 0;
}