洛谷 P3806 【模板】点分治1

题目链接

题意：给定一个 $n$ 个节点的树，每条树边具有一个长度 $c$ ，有 $m$ 个询问，每次询问包括一个 $k$ ，问树上是否存在一个距离为 $k$ 的点对。

思路：O(-1)

点分治。
对于一棵有根树，该树上节点的距离可以分为两种：

1. 经过该树的根，此时两点的距离即为两点到根的距离相加
2. 没有经过该树的根

第二种情况下的两点，其实必存在该树的一个子树(不只有一个)，这两个点的距离经过该子树的根，这又转化为第一种情况。

按照这个结论，我们可以先对整棵树处理第一种情况的点对，再把根节点去掉，产生的每棵子树同样再处理第一种情况，不断按照以上操作划分每棵树，这样的思路就是点分治了。

但是如果按照这个思路，在随意找一个节点为根的情况下，如果树是单链的形态，整个算法的时间复杂度可能达到 $n^2$ 。

这里复杂度主要在于每次分治出来的子树的大小，所以应该使每次划分出来的子树的最大节点数尽量小。这里就要求每棵树的重心，以树的重心为根划分，设树的大小为 $N$ ，则每次划分出来的子树大小不会超过 $\frac{N}{2}$ 。总划分次数不会超过 $logN$ ，每次划分处理的节点数总是 $N$ ，所以总的复杂度为 $NlogN$。

处理情况1的点对：
这里对每个子树进行 $dfs$ ，求每个点距离根的长度，每处理完一棵子树，枚举所有询问的长度以及该子树中存在的所有长度，在 $judge$ 数组中(即已经处理过的该树的子树中存在的路径长度)寻找是否有匹配的长度，之后把该子树中存在的路径长度加入一个判断数组 $judge$ 中。

PS：处理完整棵树之后要把判断数组清空(用 $memset$ 大概率会超时)。

#include
#include
#include
#include
#include
#include
#include
using namespace std;
#define ffor(i,d,u) for(int i=(d);i<=(u);++i)
#define _ffor(i,u,d) for(int i=(u);i>=(d);--i)
#define NUM 10005
#define MAXN 10000005
int n, m;
int head[NUM] = {}, edge_num = 0;
struct edge
{
    int to, next, l;
} e[NUM << 1];
bool judge[MAXN] = {}, vis[NUM] = {}, test[105] = {};//vis数组用来把根节点除去，划分当前的树
int save[NUM], leng[NUM], query[105], root;
int tot, num[NUM], ss[NUM], min_max_len;
template <typename T>
void read(T& x)
{
    x=0;
    char c;T t=1;
    while(((c=getchar())<'0'||c>'9')&&c!='-');
    if(c=='-'){t=-1;c=getchar();}
    do(x*=10)+=(c-'0');while((c=getchar())>='0'&&c<='9');
    x*=t;
}
template <typename T>
void write(T x)
{
    int len=0;char c[21];
    if(x<0)putchar('-'),x*=(-1);
    do{++len;c[len]=(x%10)+'0';}while(x/=10);
    _ffor(i, len, 1) putchar(c[i]);
}
void get_root(const int &vertex, const int &pre)//求树的重心
{
    int x, max_len = 0;
    num[vertex] = 1;
    for (int i = head[vertex]; i; i = e[i].next)
    {
        x = e[i].to;
        if (x == pre || vis[x])
            continue;
        get_root(x, vertex);
        num[vertex] += num[x];
        max_len = max(max_len, num[x]);
    }
    max_len = max(max_len, tot - num[vertex]);
    if (max_len < min_max_len)
        root = vertex, min_max_len = max_len;
}
void get_length(const int &vertex, const int &pre)//处理一棵子树里所有节点距离根节点的长度，同时更新每棵子树的大小
{
    int x;
    save[++save[0]] = leng[vertex], num[vertex] = 1;
    for (int i = head[vertex]; i; i = e[i].next)
    {
        x = e[i].to;
        if (x == pre || vis[x])
            continue;
        leng[x] = leng[vertex] + e[i].l;
        get_length(x, vertex);
        num[vertex] += num[x];
    }
}
inline void calc(const int &vertex)//对每棵子树处理距离经过根的点对
{
    int x, sum = 0;
    num[vertex] = 0;
    for (int i = head[vertex]; i; i = e[i].next)
    {
        x = e[i].to;
        if (vis[x])
            continue;
        leng[x] = e[i].l;
        get_length(x, vertex);
        num[vertex] += num[x];
        ffor(j, 1, save[0]) ffor(k, 1, m) if (query[k] >= save[j]) test[k] |= judge[query[k] - save[j]];//判断judge数组是否有匹配长度（即答案路径是否存在）
        do
            ss[++sum] = save[save[0]], judge[save[save[0]]] = true;
        while (--save[0]);//把当前子树加入judge数组中
    }
    ffor(i, 1, sum) judge[ss[i]] = false;//清空
}
void solve(const int &vertex)//点分治
{
    int x;
    judge[0] = vis[vertex] = true, calc(vertex);
    for (int i = head[vertex]; i; i = e[i].next)
    {
        x = e[i].to;
        if (vis[x])
            continue;
        tot = num[x], root = 0, min_max_len = n;
        get_root(x, 0), solve(root);
    }
}
inline void add(const int &x, const int &y, const int &length)
{
    e[++edge_num] = edge{y, head[x], length}, head[x] = edge_num;
}
inline void AC()
{
    int x, y, l;
    read(n), read(m);
    ffor(i, 2, n)
    {
        read(x), read(y), read(l);
        add(x, y, l), add(y, x, l);
    }
    ffor(i, 1, m)read(query[i]);
    tot = n, save[0] = 0, min_max_len = n;
    get_root(1, 0);
    solve(root);
    ffor(i, 1, m) puts(test[i] ? "AYE" : "NAY");
}
int main()
{
    AC();
    return 0;
}