洛谷 P3806 【模板】点分治1

typedef long long LL;
typedef pair<LL, LL> PLL;
 
const int maxm = 1e4+5;
const int maxn = 1e7+5;

int ans[maxn];
int head[maxm], cnt, siz[maxm], mxson[maxm], mxsum, rootsum, root, dis[maxm], point;
bool vis[maxm];

struct Node {
    int v, next, val;
} Nodes[maxm*2];

void addedge(int u, int v, int val) {
    Nodes[++cnt].v = v;
    Nodes[cnt].val = val;
    Nodes[cnt].next = head[u];
    head[u] = cnt;
}

void getroot(int u, int fa) {
    siz[u] = 1, mxson[u] = 0;
    for(int i = head[u]; i; i = Nodes[i].next) {
        int v = Nodes[i].v;
        if(v == fa || vis[v]) continue;
        getroot(v, u);
        siz[u] += siz[v];
        mxson[u] = max(mxson[u], siz[v]);
    }
    mxson[u] = max(mxson[u], rootsum - siz[u]);
    if(mxson[u] < mxsum) {
        root = u;
        mxsum = mxson[u];
    }
}

void getdist(int rt, int fa, int dist) {
    dis[++point] = dist;
    for(int i = head[rt]; i; i = Nodes[i].next) {
        int v = Nodes[i].v;
        if(v == fa || vis[v]) continue;
        getdist(v, rt, dist+Nodes[i].val);
    }
}

void solve(int rt, int dist, int key) {
    point = 0;
    getdist(rt, 0, dist);
    for(int i = 1; i < point; ++i)
        for(int j = i+1; j <= point; ++j)
            ans[dis[i]+dis[j]] += key;
}

void Divide(int rt) {
    vis[rt] = true;
    solve(rt, 0, 1);
    for(int i = head[rt]; i; i = Nodes[i].next) {
        int v = Nodes[i].v;
        if(vis[v]) continue;
        solve(v, Nodes[i].val, -1);
        rootsum = siz[v], mxsum = 0x3f3f3f3f;
        root = 0;
        getroot(v, 0);Divide(root);
    }
}

int main() {
    ios::sync_with_stdio(false), cin.tie(0);
    int n, m, u, v, c, K;
    cin >> n >> m;
    for(int i = 0; i < n-1; ++i) {
        cin >> u >> v >> c;
        addedge(u, v, c), addedge(v, u, c);
    }
    mxsum = 0x3f3f3f3f, rootsum = n;
    root = 0, getroot(1, 0);
    Divide(root);
    for(int i = 0; i < m; ++i) {
        cin >> K;
        if(ans[K]) cout << "AYE
";
        else cout << "NAY
";
    }
    return 0;
}