树分治 点分治poj 2114

存在2点间距离==k 输出AYE

否则输出NAY

#include<stdio.h>
#include<string.h>
#include<algorithm>
#include<vector>

using namespace std;
#define MAXN 100010
int n,m;
int cnt,ans,mi,root,k;
int head[MAXN],mx[MAXN],size[MAXN];
bool vis[MAXN];

struct edg
{
    int w,v,next;
}x[MAXN];

void add(int u,int v,int w)
{
    x[cnt].next=head[u];
    x[cnt].v=v;
    x[cnt].w=w;
    head[u]=cnt++;
}

void dfssize(int u,int fa)
{
    size[u]=1;
    mx[u]=0;
    for(int i=head[u];i!=-1;i=x[i].next)
    {
        int v=x[i].v;
        if(v!=fa&&!vis[v])
        {
            dfssize(v,u);
            size[u]+=size[v];
            if(size[v]>mx[u])
                mx[u]=size[v];
        }
    }

}
void dfsroot(int r,int u,int fa)
{
    if(size[r]-size[u]>mx[u])
        mx[u]=size[r]-size[u];
    if(mi>mx[u])
    {
        mi=mx[u];
        root=u;
    }
    for(int i=head[u];i!=-1;i=x[i].next)
    {
        int v=x[i].v;
        if(!vis[v]&&v!=fa)
            dfsroot(r,v,u);
    }
}
int num;
int dis[MAXN];

void dfsdis(int u,int d,int fa)
{
    dis[num++]=d;
    for(int i=head[u];i!=-1;i=x[i].next)
    {
        int v=x[i].v;
        if(!vis[v]&&v!=fa)
            dfsdis(v,d+x[i].w,u);
    }
}
int calc(int u,int d)
{
    num=0;
    int ret=0;
    dfsdis(u,d,0);
    sort(dis,dis+num);
    int i=0,j=num-1;
    while(i<j)//就这边该一下
    {
        if(dis[i]+dis[j]<k)
            i++;
        else if(dis[i]+dis[j]>k)
            j--;
        else
        {
            if(dis[i]==dis[j])
            {
                ret+=(j-i+1)*(j-i)/2;
                break;
            }
            int st=i,en=j;
            while(dis[st]==dis[i])
                st++;
            while(dis[en]==dis[j])
                en--;
            ret+=(st-i)*(j-en);
            i=st;
            j=en;
        }
    }
    return ret;
}
void dfs(int u)
{
    mi=n;
    dfssize(u,0);
    dfsroot(u,u,0);
    ans+=calc(root,0);
    vis[root]=1;
    for(int i=head[root];i!=-1;i=x[i].next)
    {
        int v=x[i].v;
        if(!vis[v])
        {
            ans-=calc(v,x[i].w);
            dfs(v);
        }
    }
}
int main()
{
        while(scanf("%d",&n)!=EOF&&n)
        {
            cnt=0;
            memset(head,-1,sizeof(head));

            for(int i=1;i<=n;i++)
            {
                int u,v,w;
                while(scanf("%d",&v)!=EOF&&v)
                {
                    scanf("%d",&w);
                    add(i,v,w);
                    add(v,i,w);
                }
            }
            while(scanf("%d",&k)!=EOF&&k)
            {
                ans=0;
                memset(vis,0,sizeof(vis));
                dfs(1);
                if(ans>0)
                    printf("AYE
");
                else
                    printf("NAY
");
            }
            printf(".
");
        }
        return 0;
}