poj2114 树分治（点分治）

Posted 2021-01-15 zsben991126

poj1741板子套一套，统计对数的方式改一下，可以在O(n)时间内统计对数

最后不要忘记输出最后的“.”

/*
给定一棵边权树，是否存在一条路径使得其长度为恰好为x 
把1741的板子改为求点对之间的距离=k的点对数即可 
*/
#include<iostream>
#include<cstring>
#include<cstdio>
#include<algorithm>
#define MAXN 10010
using namespace std;
int N,K;
int ans,root,Max;
struct node{
    int v,next,w;
}edge[MAXN<<1];
int head[MAXN],tot;
int size[MAXN];
int maxv[MAXN];
int vis[MAXN];
int dis[MAXN];
int num;
void init(){
    tot=ans=0;
    memset(head,-1,sizeof head);
    memset(vis,0,sizeof vis);
}
void addedge(int u,int v,int w){
    edge[tot].v=v;
    edge[tot].w=w;
    edge[tot].next=head[u];
    head[u]=tot++;
}
//一次dfs处理子树的大小 
void dfssize(int u,int f){
    size[u]=1;
    maxv[u]=0;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v==f||vis[v]) continue;
        dfssize(v,u);
        size[u]+=size[v];
        if(size[v]>maxv[u])maxv[u]=size[v];
    }
}
//一次dfs找重心，这里的r不是重心 
void dfsroot(int r,int u,int f){
    if(size[r]-size[u]>maxv[u])
        maxv[u]=size[r]-size[u];
    if(maxv[u]<Max) Max=maxv[u],root=u;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v==f||vis[v]) 
            continue;
        dfsroot(r,v,u);
    } 
}
//一次dfs求每个点到重心的距离
void dfsdis(int u,int d,int f){
    dis[num++]=d;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(v!=f && !vis[v])
            dfsdis(v,d+edge[i].w,u);
    }
}
int calc(int u,int d){
    int ret=0;
    num=0;
    dfsdis(u,d,0);//求每个点到根的距离 
    sort(dis,dis+num);
    for (int l=0, r=num-1; l<r; ) {
        if (dis[l] + dis[r] == K) {
            if (dis[l] == dis[r]) {
                ret += (r-l+1)*(r-l)/2; break;
            }
            int i=l, j=r;
            while (dis[i] == dis[l]) i++;
            while (dis[j] == dis[r]) j--;
            ret += (i-l)*(r-j);
            l = i, r = j;
        } else if (dis[l] + dis[r] < K) l++;
        else r--;
    }
    return ret;
}
void dfs(int u){//注意这里的u并不是重心 
    Max=N;
    dfssize(u,0);//求每个子树的规模 
    dfsroot(u,u,0);//求重心 
    ans+=calc(root,0);//求以root为根，dis[u]+dis[v]的点对有多少 
    vis[root]=1;//把这个点从点集删掉 
    for(int i=head[root];i!=-1;i=edge[i].next){
        int v=edge[i].v;
        if(!vis[v]){
            ans-=calc(v,edge[i].w);
            dfs(v);
        }
    } 
}
int main(){
    while(scanf("%d",&N) && N){
        init();
        int v,w;
        for(int i=1;i<=N;i++){
            while(scanf("%d",&v) && v!=0){
                scanf("%d",&w);
                addedge(i,v,w);
                addedge(v,i,w);
            }
        }
        while(scanf("%d",&K) && K!=0){
            memset(vis,0,sizeof vis);
            ans=0;
            dfs(1);
            if(ans) puts("AYE");
            else puts("NAY");
        }
        puts(".");
    }
    return 0;
}