洛谷 P3806 【模板】点分治1-树分治(点分治，容斥版) 模板题-树上距离为k的点对是否存在

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<=10000,K<=10000000

直接按照POJ 1741的代码改的，其他的没什么。POJ 1741.Tree-树分治(点分治) 模板题-区间点对最短距离<=K的点对数量

代码：

//树分治-点分治
#include<bits/stdc++.h>
using namespace std;
const int inf=1e9+;
const int maxn=1e4+;
const int maxm=1e7+;

int head[maxn<<],tot;
int root,allnode,n,m,k;
int vis[maxn],deep[maxn],dis[maxn],siz[maxn],maxv[maxn];//deep[0]子节点个数(路径长度)，maxv为重心节点
int ans[maxm];

struct node{
    int to,next,val;
}edge[maxn<<];

void add(int u,int v,int w)//前向星存图
{
    edge[tot].to=v;
    edge[tot].next=head[u];
    edge[tot].val=w;
    head[u]=tot++;
}

void init()//初始化
{
    memset(head,-,sizeof head);
    memset(vis,,sizeof vis);
    tot=;
}

void get_root(int u,int father)//重心
{
    siz[u]=;maxv[u]=;
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        if(v==father||vis[v]) continue;
        get_root(v,u);//递归得到子树大小
        siz[u]+=siz[v];
        maxv[u]=max(maxv[u],siz[v]);//更新u节点的maxv
    }
    maxv[u]=max(maxv[u],allnode-siz[u]);//保存节点size
    if(maxv[u]<maxv[root]) root=u;//更新当前子树的重心
}

void get_dis(int u,int father)//获取子树所有节点与根的距离
{
    deep[++deep[]]=dis[u];
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        if(v==father||vis[v]) continue;
        int w=edge[i].val;
        dis[v]=dis[u]+w;
        get_dis(v,u);
    }
}

void cal(int u,int now,int val)
{
    dis[u]=now;deep[]=;
    get_dis(u,);
    sort(deep+,deep+deep[]+);
    for(int i=;i<=deep[];i++){
        for(int j=;j<=deep[];j++){
            if(i!=j) ans[deep[i]+deep[j]]+=val;
        }
    }
}

void solve(int u)
{
    cal(u,,);
    vis[u]=;
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        int w=edge[i].val;
        if(vis[v]) continue;
        cal(v,w,-);
        allnode=siz[v];
        root=;
        get_root(v,u);
        solve(root);
    }
}

int main()
{
    scanf("%d%d",&n,&m);
    init();
    for(int i=;i<n;i++){
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        add(u,v,w);
        add(v,u,w);
    }
    root=;allnode=n;maxv[]=inf;
    get_root(,);
    solve(root);
    while(m--){
        scanf("%d",&k);
        if(ans[k]){
            printf("AYE\n");
        }
        else{
            printf("NAY\n");
        }
    }
    return ;
}