题目描述

给定一棵有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

// luogu-judger-enable-o2
#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
const int M = 10005, K = 10000001;

int n, tot, cnt, Ans, root, sum, h[M], siz[M], dis[M], tong[M], f[M], num[K+1];
bool vis[M];
struct edge{int v, nxt, w;}G[M<<1];
void add(int u, int v, int w){
    G[++tot].v = v; G[tot].w = w; G[tot].nxt = h[u]; h[u] = tot;
}
void init(){
    memset(vis, 0, sizeof(vis));
    memset(h, 0, sizeof(h));
    tot = 0; Ans = 0; sum = 0;
}


void getdeep(int u, int fa){
    tong[++cnt] = dis[u];
    for(int i = h[u]; i; i = G[i].nxt){
        int v = G[i].v;
        if(vis[v] || v == fa)continue;
        dis[v] = dis[u] + G[i].w;    
        getdeep(v, u);    
    }
}

void getroot(int u, int fa){
    siz[u] = 1; f[u] = 0;
    for(int i = h[u]; i; i = G[i].nxt){
        int v = G[i].v;
        if(vis[v] || v == fa)continue;
        getroot(v, u);
        siz[u] += siz[v];
        f[u] = max(f[u], siz[v]);
    }
    f[u] = max(f[u], sum - siz[u]);
    if(f[u] < f[root]) root = u; 
    
}
void cal(int u, int ret, bool opt){
    dis[u] = ret;
    cnt = 0; 
    getdeep(u, 0);
    int ans = 0;
    for(int i = 1; i <= cnt; i++)
        for(int j = i + 1; j <= cnt; j++){
            if(opt && tong[i] + tong[j] < K) ++num[tong[i] + tong[j]];
            else if(tong[i] + tong[j] < K) --num[tong[i] + tong[j]];    
        }
}

void solve (int u){
    cal(u, 0, 1);
    vis[u] = 1;
    for(int i = h[u]; i; i = G[i].nxt){
        int v = G[i].v;
        if(vis[v])continue;
        cal(v, G[i].w, 0);
        sum = siz[v];
        getroot(v, root = 0);
        solve(root);
    }
}


int main(){
    f[0] = 1e9;
    int m, fg = 0;
    scanf("%d%d", &n, &m);
    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);
        if(!w)fg = 1;
    }
        
    solve(1);
    num[0] = fg;
    while(m--){
        int k; scanf("%d", &k);
        if(num[k])puts("AYE");
        else puts("NAY");
    }
}

View Code

P4178 Tree

链接：https://www.luogu.org/problemnew/show/P4178