P3806 【模板】点分治1 | 点分治模板题

点分治模板题

题意：
给定一棵有 n 个点的树。询问树上距离为 k 的点对是否存在。
分析：
点分治模板
代码：

//copy AgOH's code
#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e5 + 10;
const int maxnum = 1e8 + 10;
struct edge
{
    int ep, v, nex;
    edge(int a = 0, int b = 0, int c = 0) : ep(a), v(b), nex(c){}
}e[maxn << 1];
int head[maxn], tot;
int n, m;

void init() {memset(head, -1, sizeof(head)); tot = 0;}
void addedge(int sp, int ep, int v)
{
    e[tot] = edge(ep, v, head[sp]);
    head[sp] = tot++;
    e[tot] = edge(sp, v, head[ep]);
    head[ep] = tot++;
}

int maxp[maxn], siz[maxn], vis[maxn], treesiz, root;
int dis[maxn], judge[maxnum], tmp[maxn], tmpnum;
int q[maxn], ok[maxn];

void getroot(int now, int f)
{
    siz[now] = 1, maxp[now] = 0;
    for(int i = head[now]; ~i; i = e[i].nex)
    {
        int ep = e[i].ep;
        if(ep == f || vis[ep]) continue;
        getroot(ep, now);
        siz[now] += siz[ep];
        maxp[now] = max(maxp[now], siz[ep]);
    }
    maxp[now] = max(maxp[now], treesiz - siz[now]);
    if(maxp[now] < maxp[root]) root = now;
}

void getdis(int now, int f)
{
    tmp[tmpnum++] = dis[now];
    for(int i = head[now]; ~i; i = e[i].nex)
    {
        int ep = e[i].ep, v = e[i].v;
        if(vis[ep] || ep == f) continue;
        dis[ep] = dis[now] + v;
        getdis(ep, now);
    }
}

void solve(int now)
{
    queue<int> que;
    for(int i = head[now]; ~i; i = e[i].nex)
    {
        int ep = e[i].ep, v = e[i].v;
        if(vis[ep]) continue;
        tmpnum = 0;
        dis[ep] = v;
        getdis(ep, now);
        for(int j = 0; j < tmpnum; j++)
            for(int k = 0; k < m; k++)
                if(q[k] >= tmp[j])
                    ok[k] |= judge[q[k] - tmp[j]];
        for(int j = 0; j < tmpnum; j++)
        {
            que.push(tmp[j]);
            judge[tmp[j]] = true;
        }
    }
    while(!que.empty())
    {
        judge[que.front()] = false;
        que.pop();
    }
}

void divide(int now)
{
    vis[now] = judge[0] = true;
    solve(now);
    for(int i = head[now]; ~i; i = e[i].nex)
    {
        int ep = e[i].ep;
        if(vis[ep]) continue;
        maxp[root = 0] = treesiz = siz[ep];
        getroot(ep, -1);
        getroot(root, -1);
        divide(root);
    }
}

int main()
{
    cin >> n >> m;
    init();
    for(int i = 1; i < n; i++)
    {
        int x, y, v;
        scanf("%d%d%d", &x, &y, &v);
        addedge(x, y, v);
    }
    for(int i = 0; i < m; i++)
        scanf("%d", q + i);
    maxp[root = 0] = treesiz = n;
    getroot(1, -1);
    getroot(root, -1);
    divide(root);
    for(int i = 0; i < m; i++)
    {
        if(ok[i]) puts("AYE");
        else puts("NAY");
    }
}