洛谷P3806 【模板】点分治 1 题解

题意：

给定一棵有 \(n\) 个点的树，询问树上距离为 \(k\) 的点对是否存在。

输入格式：

第一行两个数 \(n,m\)。

第 \(2\) 到第 \(n\) 行，每行三个整数 \(u, v, w\)，代表树上存在一条连接 \(u\) 和 \(v\) 边权为 \(w\) 的路径。

接下来 \(m\) 行，每行一个整数 \(k\)，代表一次询问。

输出格式：

对于每次询问输出一行一个字符串代表答案，存在输出 AYE，否则输出 NAY。

数据范围：

保证 \(1 \leq n\leq 10^4\)，\(1 \leq m\leq 100\)，\(1 \leq k \leq 10^7\)，\(1 \leq u, v \leq n\)，\(1 \leq w \leq 10^4\)。

本题是点分治的模板题。点分治的算法思想很简单。

考虑点 \(u\) 作为根节点时，任意两个点产生的路径可以被分为：经过 \(u\) 的、不经过 \(u\) 的。

前者有 \(d(i,j) =d(i,u) + d(u,j)\) ，而后者并不是很好求解。

点分治给出的解法是，对于 \(u\) 的子节点，重复这个“将某个点作为根”的操作

可以发现，如果不断递归下去，可以把所有的路径都统计出来，复杂度是 \(\mathcal{O}(nd)\) ，\(d\) 为递归层数

显然我们要减少这个递归层数，这个很简单，只要取树的重心即可。

算法思想固然简单，实现的时候细节还是挺多的，比如这个根节点的父节点需要特判一下。

时间复杂度 \(\mathcal{O}(n \log n+nm)\)

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
#define N ((int)(1e4 + 15))
#define M ((int)(115))
#define K ((int)(1e7 + 15))

int n, m, rt, rt_fa, sum, pos = 1, head[N], sz[N], mx[N], qry[N];
int tot, q[N], val[N], dis[N]; char ans[M], vis[N], pd[K];
struct Edge { int u, v, w, next; } e[N * 2];
void addEdge(int u, int v, int w) {
    e[++pos] = {u, v, w, head[u]}; head[u] = pos;  
}
void getroot(int u, int fa)
{
    mx[u] = 0; sz[u] = 1;
    for(int i = head[u]; i; i = e[i].next)
    {
        int v = e[i].v; if(v == fa || vis[v]) continue;
        getroot(v, u); sz[u] += sz[v]; up(mx[u], sz[v]);
    }
    up(mx[u], sum - sz[u]); if(mx[u] < mx[rt]) { rt = u, rt_fa = fa; }
}
void getdis(int u, int fa)
{
    if(dis[u] > 10000000) return; 
    val[++tot] = dis[u];
    for(int i = head[u]; i; i = e[i].next)
    {
        int v = e[i].v; if(v == fa || vis[v]) continue;
        dis[v] = dis[u] + e[i].w; getdis(v, u);
    }
}
void init(int u)
{
    int c = 0;
    for(int i = head[u]; i; i = e[i].next)
    {
        int v = e[i].v; if(vis[v]) continue;
        tot = 0; dis[v] = e[i].w; getdis(v, u);
        for(int j = 1; j <= tot; j++)
            for(int k = 1; k <= m; k++)
                if(qry[k] >= val[j]) ans[k] |= pd[qry[k] - val[j]];
        for(int j = 1; j <= tot; j++) { q[++c] = val[j], pd[val[j]] = true; }
    }
    for(int i = 1; i <= c; i++) pd[q[i]] = false;
}
void solve(int u)
{
    vis[u] = pd[0] = true; init(u);
    for(int i = head[u]; i; i = e[i].next)
    {
        int v = e[i].v; if(vis[v]) continue;
        sum = v == rt_fa ? n - sz[u] : sz[v];
        mx[rt = rt_fa = 0] = n; getroot(v, u); solve(rt);
    }
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n >> m; 
    for(int i = 1, u, v, w; i < n; i++)
    {
        cin >> u >> v >> w;
        addEdge(u, v, w); addEdge(v, u, w);
    }
    for(int i = 1; i <= m; i++) cin >> qry[i];
    mx[rt = rt_fa = 0] = sum = n;
    getroot(1, 0); solve(rt);
    for(int i = 1; i <= m; i++) cout << (ans[i] ? "AYE" : "NAY") << '\n';
    return 0;
}