poj2114 Boatherds

题目：给你一棵n个节点的树，每条边都有一个权值，再给你一个K，问你是否存在两个点u,v，使得uv之间的距离之和为K。

思路：点分治

代码：

#pragma comment(linker, "/STACK:1024000000,1024000000")
#include<iostream>
#include<algorithm>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<string>
#include<vector>
#include<map>
#include<set>
#include<queue>
#include<stack>
#include<list>
#include<numeric>
using namespace std;
#define PI acos(-1.0)
#define LL long long
#define ULL unsigned long long
#define INF 0x3f3f3f3f
#define mm(a,b) memset(a,b,sizeof(a))
#define PP puts("*********************");
template<class T> T f_abs(T a){ return a > 0 ? a : -a; }
template<class T> T gcd(T a, T b){ return b ? gcd(b, a%b) : a; }
template<class T> T lcm(T a,T b){return a/gcd(a,b)*b;}
// 0x3f3f3f3f3f3f3f3f
// 0x3f3f3f3f

const int maxn=5e4+50;
struct Node{
    int v,w,next;
}edge[2*maxn];
int head[maxn],tol;//节点从1开始计数
int son[maxn],f[maxn],vis[maxn];
//son表示以i为根的子树大小
//f数组表示以u为根的最大子树的大小
//vis表示该节点是不是已经作为重心了
int dep[maxn],siz,d[maxn];
int cntv,root,K,ans;
void init(){
    mm(head,-1);tol=0;
    mm(vis,0);
    f[0]=INF;//一定要初始化成一个很大的值
}
void addedge(int u,int v,int w){
    edge[tol].v=v;
    edge[tol].w=w;
    edge[tol].next=head[u];
    head[u]=tol++;
}
void getroot(int u,int fa){//寻找重心
    son[u]=1,f[u]=0;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v==fa||vis[v]) continue;
        getroot(v,u);
        son[u]+=son[v];
        f[u]=max(f[u],son[v]);
    }
    f[u]=max(f[u],cntv-son[u]);//sum表示当前树的大小
    if(f[u]<f[root]) root=u;//更新当前重心
}
void getdepth(int u,int fa){
    dep[siz++]=d[u];
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v==fa||vis[v]) continue;
        d[v]=d[u]+edge[i].w;
        getdepth(v,u);
    }
}
int cal(int u,int w){
    d[u]=w;
    siz=0;
    getdepth(u,0);
    sort(dep,dep+siz);
    int l=0,r=siz-1,res=0;
    while(l<r){
        if(dep[l]+dep[r]<K)
            l++;
        else if(dep[l]+dep[r]>K)
            r--;
        else{
            if(dep[l]==dep[r]){
                res+=(r-l+1)*(r-l)/2;
                break;
            }
            else{
                int s=l,t=r;
                int x=1,y=1;
                while(dep[s]==dep[s+1]){
                    x++;
                    s++;
                }
                while(dep[t]==dep[t-1]){
                    y++;
                    t--;
                }
                res+=x*y;
                l=s+1;
                r=t-1;
            }
        }
    }
    return res;
}
void solve(int u){//计算以u为重心的树
    ans+=cal(u,0);
    vis[u]=1;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(vis[v]) continue;
        ans-=cal(v,edge[i].w);//减去不合法的数目
        cntv=son[v];
        root=0;
        getroot(v,0);
        solve(root);
    }
}
int main(){

    int n,u,v,w;
    while(~scanf("%d",&n)){
        if(!n) break;
        init();
        for(int i=1;i<=n;i++){
            while(~scanf("%d",&v)){
                if(!v) break;
                scanf("%d",&w);
                addedge(i,v,w);
                addedge(v,i,w);
            }
        }
        while(~scanf("%d",&K)){
            if(!K) break;
            mm(vis,0);
            cntv=n;
            root=0;
            getroot(1,0);
            ans=0;
            solve(root);
            if(ans>0) printf("AYE\n");
            else printf("NAY\n");
        }
        printf(".\n");
    }
    return 0;
}