洛谷 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\n");
        }
        else{
            printf("NAY\n");
        }
    }
    return 0;
}