洛谷P3806 【模板】点分治1

【模板】点分治1

题目传送门

代码如下

#include <bits/stdc++.h>
#define INF 0x3f3f3f3f
using namespace std;
typedef long long ll;

const int N = 10005;
struct T{
    int r, w;
    T(int r, int w): r(r), w(w){}
};
vector<T> g[N];

int rt, num, n, m;;
int siz[N];
bool vis[N];

void find_root(int x, int fa)
{
    siz[x] = 1;
    int max_son = 0;
    for(int i = 0; i < g[x].size(); i ++){
        int r = g[x][i].r;
        if(vis[r] || r == fa) continue;
        find_root(r, x);
        siz[x] += siz[r];
        max_son = max(max_son, siz[r]);
    }
    max_son = max(max_son, n - siz[x]);
    if(max_son < num){
        num = max_son;
        rt = x;
    }
}

int d[N], b[N];
vector<int> vec;

void get_dis(int x, int fa, int s)
{
    for(int i = 0; i < g[x].size(); i ++){
        int r = g[x][i].r;
        if(vis[r] || r == fa) continue;
        b[r] = s, d[r] = d[x] + g[x][i].w;
        vec.push_back(r);
        get_dis(r, x, s);
    }
}

bool cmp(int a, int b)
{
    return d[a] < d[b];
}

bool calc(int k)
{
    sort(vec.begin(), vec.end(), cmp);
    int l, r;
    l = 0, r = vec.size() - 1;
    while(l < r){
        int x = vec[l], y = vec[r];
        if(d[x] + d[y] > k){
            --r;
        }
        else {
            if(d[x] + d[y] == k && b[x] != b[y])
                return true;
            ++l;
        }
    }
    return false;
}

bool dfs(int x, int k)
{
    num = n;
    find_root(x, 0);
    int t = rt;
    vis[t] = true;
    vec.clear();
    b[t] = d[t] = 0;
    vec.push_back(t);
    for(int i = 0; i < g[t].size(); i ++){
        int r = g[t][i].r;
        if(!vis[r]){
            b[r] = r, d[r] = g[t][i].w;
            vec.push_back(r);
            get_dis(r, t, r);
        }
    }
    if(calc(k))
        return true;
    for(int i = 0; i < g[t].size(); i ++){
        int r = g[t][i].r;
        if(vis[r]) continue;
        if(dfs(r, k)) return true;
    }
    return false;
}

int main()
{
    scanf("%d%d", &n, &m);
    for(int i = 1; i < n; i ++){
        int a, b, c;
        scanf("%d%d%d", &a, &b, &c);
        g[a].push_back(T(b, c));
        g[b].push_back(T(a, c));
    }
    for(int i = 1; i <= m; i ++){
        int k;
        scanf("%d", &k);
        memset(vis, 0, sizeof(vis));
        if(dfs(1, k)){
            puts("AYE");
        }
        else {
            puts("NAY");
        }
    }
    return 0;
}