算法竞赛进阶指南——0x45【点分治】

树上淀粉质点分治

P3806 【模板】点分治1

#include <bits/stdc++.h>
// #pragma GCC optimize(3, "Ofast", "inline")
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<long long, long long> pll;
#define IOS ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
const int N = 1e5 + 10;
const int M = 1e7 + 1000;
const int mod = 11092019;
const int inf = 0x7f7f7f7f;
const int eps = 1e-8;
ll read()
{
    ll x = 0, f = 1;
    char c = getchar();
    while (c < '0' || c > '9')
    {
        if (c == '-')
            f = -1;
        c = getchar();
    }
    while ('0' <= c && c <= '9')
    {
        x = x * 10 + c - '0';
        c = getchar();
    }
    return x * f;
}
ll n, m, tot, root, sumsiz;
ll head[N], to[N << 1], nex[N << 1], val[N << 1], que[110], maxsiz[N], siz[N], dis[N], notedis[N], notejud[N];
bool jud[M], ans[110], vis[N];
void addedge(ll u, ll v, ll w)
{
    to[++tot] = v;
    val[tot] = w;
    nex[tot] = head[u];
    head[u] = tot;
}
void getroot(ll rt, ll fa)
{ //简单的找重心
    siz[rt] = 1, maxsiz[rt] = 0;
    for (int i = head[rt]; i; i = nex[i])
    {
        if (vis[to[i]] || to[i] == fa)
            continue; //加了一个vis判断，防止跑到已经访问过的根节点给上去。
        getroot(to[i], rt);
        siz[rt] += siz[to[i]];
        maxsiz[rt] = max(maxsiz[rt], siz[to[i]]);
    }
    maxsiz[rt] = max(maxsiz[rt], sumsiz - siz[rt]); //
    if (maxsiz[root] > maxsiz[rt])
        root = rt;
}
void getdis(ll rt, ll fa)
{                                    //就是dfs树上最短路的实现过程
    notedis[++notedis[0]] = dis[rt]; //记录其子树的每个节点到根节点的距离
    for (int i = head[rt]; i; i = nex[i])
    {
        if (vis[to[i]] || to[i] == fa)
            continue;
        dis[to[i]] = dis[rt] + val[i];
        getdis(to[i], rt);
    }
}
void calc(ll rt)
{ //核心
    ll cnt = 0;
    for (int i = head[rt]; i; i = nex[i]) //
    {
        if (vis[to[i]])
            continue;        //同样的也是访问子树
        dis[to[i]] = val[i]; //这里一定要记得重置
        notedis[0] = 0;
        getdis(to[i], rt);
        for (int j = 1; j <= notedis[0]; j++) //查询有没有点到当前子树的点的距离是符合que中的要求的
            for (int k = 1; k <= m; k++)
                if (que[k] >= notedis[j])               //
                    ans[k] |= jud[que[k] - notedis[j]]; //
        for (int j = 1; j <= notedis[0]; j++)           //记录我们jud中被标记的点，方便在下一次分治之前重置
            if (notedis[j] <= 1e7 + 5)                  //特判一下吧，题目的dis可能会到1e8，为了防止数组越界
            {
                notejud[++cnt] = notedis[j];
                jud[notedis[j]] = 1;
            }
    }
    for (int i = 1; i <= cnt; i++) //不用memset重置，防止变成n^2的算法
        jud[notejud[i]] = 0;
}
void solve(ll rt)
{
    vis[rt] = jud[0] = 1; //这个点被访问过，防止其子树上的点再次访问这个点
    calc(rt);
    for (int i = head[rt]; i; i = nex[i])
    {
        if (vis[to[i]])
            continue; //我们肯定是找一个没有访问的子树上的点去进行下一次分治递归
        sumsiz = siz[to[i]];
        root = 0;
        maxsiz[root] = inf;
        getroot(to[i], 0);
        solve(root);
    }
}
int main()
{
    // IOS;
    cin >> n >> m;
    for (int i = 1; i < n; i++) //
    {
        int x, y, z;
        cin >> x >> y >> z;
        addedge(x, y, z);
        addedge(y, x, z);
    }
    for (int i = 1; i <= m; i++)
        cin >> que[i];
    root = 0; //寻找初始的递归根节点
    maxsiz[root] = inf;
    getroot(1, 0);
    solve(root);
    for (int i = 1; i <= m; i++)
        puts(ans[i] ? "AYE" : "NAY");
    return 0;
}

Tree