萌新点分治 TLE #7 #9 求调

https://www.luogu.com.cn/record/116895867

#include <bits/stdc++.h>
#define Foredge(i, u, v, w) for (int i = head[u], v = e[i].to, w = e[i].w; i; i = e[i].next, v = e[i].to, w = e[i].w)
using namespace std;
namespace Milkcat {
    typedef long long LL;
    const int N = 1e6 + 5;
    struct edge { LL w, next, to; } e[N];
    int n, q, u, v, w, cnt, root, Q[N], head[N], siz[N], mx[N], vis[N];
    map<int, int> ans, S;
    void add(int u, int v, LL w) { e[++ cnt] = {w, head[u], v}, head[u] = cnt; }
    void getRoot(int u, int fa, int S) {
        siz[u] = 1, mx[u] = 0;
        Foredge(i, u, v, w)
            if (v != fa && !vis[v]) getRoot(v, u, S), siz[u] += siz[v], mx[u] = max(mx[u], siz[v]);
        mx[u] = max(mx[u], S - siz[u]);
        if (!root || mx[u] < mx[root]) root = u;
    }
    void getDis(int u, int fa, int dep) {
        S[dep] ++;
        Foredge(i, u, v, w)
            if (v != fa && !vis[v]) getDis(v, u, dep + w);
    }
    void solve(int u, int len, int w) {
        S.clear(), getDis(u, 0, len);
        for (int i = 1; i <= q; i ++)
            for (auto val : S)
                if (S.count(Q[i] - val.first)) ans[i] += S[Q[i] - val.first] * val.second * w;
    }
    void divide(int u) {
        solve(u, 0, 1), vis[u] = 1;
        Foredge(i, u, v, w)
            if (!vis[v]) solve(v, w, -1), root = 0, getRoot(v, 0, siz[v]), divide(root);
    }
    int main() {
        cin >> n >> q;
        for (int i = 1; i < n; i ++)
            cin >> u >> v >> w, add(u, v, w), add(v, u, w);
        for (int i = 1; i <= q ; i ++) cin >> Q[i];
        getRoot(1, 0, n), divide(root);
        for (int i = 1; i <= q; i ++)
            puts(ans[i] / 2 ? "AYE" : "NAY");
        return 0;
    }
}
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    int T = 1;
    while (T --) Milkcat::main();
    return 0;
}