P3806 点分治1 (点分治模板题)

题目描述

给定一棵有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<=10000,K<=10000000

思路:

模板题
预处理所有 点对 间距离 出现的次数
然后每个询问O(1)解决

代码少许注释

code：

O(n²)

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
using namespace std;
const int maxm=1e4+5;
const int inf=1e9;
const int K=1e7+5;//k的最大范围
int head[maxm],nt[maxm<<1],to[maxm<<1],w[maxm<<1],cnt;
int mark[maxm];//隔断
int sz[maxm];//以x为根的子树大小
int son[maxm];//x的最大子树大小
int size;//当前整棵树的大小
int root;//当前区域树的重心
//int mx;//当前区域最大子树大小
int d[maxm],num;//当前区域各点到重心的距离
int ans[K];//出现的答案的次数
int n,m;
void add(int x,int y,int z){
    cnt++;nt[cnt]=head[x];head[x]=cnt;to[cnt]=y;w[cnt]=z;
}
void getroot(int x,int fa){//找重心
    sz[x]=1;
    son[x]=0;
    for(int i=head[x];i;i=nt[i]){
        int v=to[i];
        if(mark[v]||v==fa)continue;
        getroot(v,x);
        sz[x]+=sz[v];
        son[x]=max(son[x],sz[v]);
    }
    son[x]=max(son[x],size-sz[x]);//因为这一步所以有可能与重心x相连的点v的sz[v]包含x
//    if(son[x]<mx){
//        mx=son[x];
//        root=x;
//    }
    if(son[root]>son[x]){
        root=x;
    }
}
void ask(int x,int val,int fa){//求当前区域内各点到重心的距离
    d[++num]=val;
    for(int i=head[x];i;i=nt[i]){
        int v=to[i];
        if(mark[v]||v==fa)continue;
        ask(v,val+w[i],x);
    }
}
void solve(int x,int val,int change){
    num=0;
    ask(x,val,0);
    for(int i=1;i<=num;i++){//枚举两点
        for(int j=i+1;j<=num;j++){
            ans[d[i]+d[j]]+=change;//这里地方有一些答案是错误的，之后要去掉(就不详细说明为什么错误了)
        }
    }
}
void divide(int x){
    mark[x]=1;//阻断x
    solve(x,0,1);
    for(int i=head[x];i;i=nt[i]){
        int v=to[i];
        if(mark[v])continue;
        solve(v,w[i],-1);//容斥，去掉错误答案
        son[root=0]=size=sz[v];
//        mx=inf;
//        size=(sz[v]>sz[x])?n-sz[x]:sz[v];//因为x可能在sz[v]中，所以size=sz[v]其实是错误的,但是也不会错
        getroot(v,0);//这两步和main里面的一样，递归处理子树
        divide(root);//
    }
}
signed main(){
    scanf("%d%d",&n,&m);
    for(int i=1;i<n;i++){//树边
        int a,b,c;
        scanf("%d%d%d",&a,&b,&c);
        add(a,b,c);
        add(b,a,c);
    }
    son[root=0]=size=n;
    getroot(1,0);//找重心
    divide(root);
    while(m--){
        int x;
        scanf("%d",&x);
        if(ans[x]){
            puts("AYE");
        }else{
            puts("NAY");
        }
    }
    return 0;
}