洛谷 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+10;
const int maxn=1e4+10;
const int maxm=1e7+10;

int head[maxn<<1],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<<1];

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,-1,sizeof head);
    memset(vis,0,sizeof vis);
    tot=0;
}

void get_root(int u,int father)//重心
{
    siz[u]=1;maxv[u]=0;
    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[0]]=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[0]=0;
    get_dis(u,0);
    sort(deep+1,deep+deep[0]+1);
    for(int i=1;i<=deep[0];i++){
        for(int j=1;j<=deep[0];j++){
            if(i!=j) ans[deep[i]+deep[j]]+=val;
        }
    }
}

void solve(int u)
{
    cal(u,0,1);
    vis[u]=1;
    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,-1);
        allnode=siz[v];
        root=0;
        get_root(v,u);
        solve(root);
    }
}

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