洛谷 P3806 【模板】点分治1

题目背景

感谢hzwer的点分治互测。

题目描述

给定一棵有n个点的树

询问树上距离为k的点对是否存在。

输入输出格式

输入格式：

n,m 接下来n-1条边a,b,c描述a到b有一条长度为c的路径

接下来m行每行询问一个K

输出格式：

对于每个K每行输出一个答案，存在输出“AYE”,否则输出”NAY”(不包含引号)

输入输出样例

输入样例#1：复制

2 1
1 2 2
2

输出样例#1：复制

AYE

说明

对于30%的数据n<=100

对于60%的数据n<=1000,m<=50

对于100%的数据n<=10000,m<=100,c<=1000,K<=10000000

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define maxn 10010
using namespace std;
int head[maxn],d[maxn],num,sum,f[maxn],n,m,tt,pp,root,son[maxn];
struct node{int to,pre,v;}e[maxn*2];
bool is[10000010],vis[maxn];
struct Node{int dis,which;}a[maxn];
void Insert(int from,int to,int v){
    e[++num].to=to;
    e[num].v=v;
    e[num].pre=head[from];
    head[from]=num;
}
void getroot(int x,int father){
    son[x]=1;f[x]=0;
    for(int i=head[x];i;i=e[i].pre){
        int to=e[i].to;
        if(to==father||vis[to])continue;
        getroot(to,x);
        son[x]+=son[to];
        f[x]=max(f[x],son[to]);
    }
    f[x]=max(f[x],sum-son[x]);
    if(f[x]<f[root])root=x;
}
Node make_Node(int x,int y){
    Node res;
    res.dis=x;res.which=y;
    return res;
}
int getdep(int r1,int r2,int x,int father){
    for(int i=head[x];i;i=e[i].pre){
        int to=e[i].to;
        if(to==father||vis[to])continue;
        if(x==r1)pp++;
        d[to]=d[x]+e[i].v;
        if(x==r1)a[++tt]=make_Node(d[to],pp);
        else a[++tt]=make_Node(d[to],r2);
        is[d[to]]=1;
        if(x==r1)getdep(r1,pp,to,x);
        else getdep(r1,r2,to,x);
    }
}
void solve(int x){
    d[x]=tt=pp=0;
    getdep(x,0,x,0);
    vis[x]=1;
    for(int i=1;i<=tt;i++)
        for(int j=i+1;j<=tt;j++)
            if(a[i].which!=a[j].which)is[a[i].dis+a[j].dis]=1;
    for(int i=head[x];i;i=e[i].pre){
        int to=e[i].to;
        if(vis[to])continue;
        sum=son[to];root=0;
        getroot(to,0);
        solve(root);
    }
}
int main(){
    scanf("%d%d",&n,&m);
    int x,y,z;
    for(int i=1;i<n;i++){
        scanf("%d%d%d",&x,&y,&z);
        Insert(x,y,z);Insert(y,x,z);
    }
    root=0;sum=f[0]=n;
    getroot(1,0);
    solve(root);
    while(m--){
        scanf("%d",&x);
        if(is[x])puts("AYE");
        else puts("NAY");
    }
    return 0;
}